%matplotlib inline
from d2l import torch as d2l
import math
import torch
from torch import nn
from torch.nn import functional as F10.2 Attention Scoring and Masking
In Section 10.1.3, distance-based kernels such as the Gaussian defined the attention weights. A more general attention layer first computes a score \(a(\mathbf{q}, \mathbf{k})\) and then normalizes the scores with softmax Equation 10.1.2. This section derives scaled dot-product attention, including the \(1/\sqrt d\) factor. Masked softmax handles padding and causal constraints, while batched matrix multiplication evaluates many queries and sequences efficiently. We conclude with the alignment model from neural machine translation that introduced learned attention.
%matplotlib inline
from d2l import jax as d2l
from flax import nnx
import jax
from jax import numpy as jnp
import math10.2.1 Dot-Product Attention
Expanding the Gaussian kernel gives the following attention score before exponentiation:
\[ a(\mathbf{q}, \mathbf{k}_i) = -\frac{1}{2} \|\mathbf{q} - \mathbf{k}_i\|^2 = \mathbf{q}^\top \mathbf{k}_i -\frac{1}{2} \|\mathbf{k}_i\|^2 -\frac{1}{2} \|\mathbf{q}\|^2. \]
First, the final term depends on \(\mathbf{q}\) only, so it is identical for all keys, and the softmax normalization Equation 10.1.2 removes it entirely. Second, if the key norms \(\|\mathbf{k}_i\|\) are all equal, the middle term drops out the same way, and the Gaussian kernel and the dot product induce identical attention weights. In general the norms are not equal, and dropping the term is a modeling decision rather than an approximation: we adopt the dot product \(\mathbf{q}^\top \mathbf{k}_i\) as a compatibility function in its own right. Learned query and key representations can shape this compatibility, while the Gaussian expansion remains an exact special case when all key norms are equal.
One adjustment controls the nominal score magnitude at initialization. Assume for this calculation that the \(d_k\) coordinates of the query and key are independent, zero-mean, unit-variance random variables. Their dot product then has zero mean and variance \(d_k\). Rescaling by \(1/\sqrt{d_k}\) removes this dimension dependence under the stated assumptions and yields the scaled dot-product attention scoring function of the Transformer (Vaswani et al. 2017):
\[ a(\mathbf{q}, \mathbf{k}_i) = \mathbf{q}^\top \mathbf{k}_i / \sqrt{d_k}. \tag{10.2.1}\]
The attention weights are obtained, as always, with the softmax:
\[\alpha(\mathbf{q}, \mathbf{k}_i) = \mathrm{softmax}(a(\mathbf{q}, \mathbf{k}_i)) = \frac{\exp(\mathbf{q}^\top \mathbf{k}_i / \sqrt{d_k})}{\sum_{j=1}^m \exp(\mathbf{q}^\top \mathbf{k}_j / \sqrt{d_k})}. \tag{10.2.2}\]
10.2.1.1 Score Variance and the \(1/\sqrt{d}\) Factor
Learned queries and keys need not retain the independence or unit-variance assumptions, so the calculation is a scaling heuristic rather than a trained-model guarantee. Its purpose is to avoid dimension-induced softmax saturation at initialization. Once one score exceeds the others by a large margin, the winning weight approaches \(1\) and the rest approach \(0\). The gradient returning along the score path shrinks with them, because the Jacobian of the softmax is
\[ \frac{\partial \boldsymbol{\alpha}}{\partial \mathbf{a}} = \mathrm{diag}(\boldsymbol{\alpha}) - \boldsymbol{\alpha} \boldsymbol{\alpha}^\top, \]
and this tends to the zero matrix as \(\boldsymbol{\alpha}\) approaches a one-hot vector, so the queries and keys behind those scores stop being updated. (Only that query–key route saturates; gradients still flow through the values, the output projection, and any residual connection.) For finite scores this Jacobian is never the zero matrix — as we noted in Section 10.1, it always keeps the all-ones vector in its null space and nothing else — but it can come arbitrarily close. We measure both effects with random queries and keys with unit-variance entries, compute attention over \(64\) candidate keys, and record the entropy of the resulting weight distribution as the dimension \(d\) grows, with and without the \(1/\sqrt{d}\) factor.
def attention_stats(d, num_keys=64, num_queries=256, scaled=True):
q = torch.randn(num_queries, d)
keys = torch.randn(num_queries, num_keys, d)
scores = (keys @ q[..., None]).squeeze(-1)
if scaled:
scores = scores / math.sqrt(d)
alpha = F.softmax(scores, dim=-1)
entropy = -(alpha * torch.log(alpha + 1e-12)).sum(-1).mean()
# Frobenius norm of the softmax Jacobian diag(alpha) - alpha alpha^T
jac = torch.diag_embed(alpha) - alpha[..., :, None] * alpha[..., None, :]
return entropy.item(), jac.square().sum((-2, -1)).sqrt().mean().item()
torch.manual_seed(0)
ds = [2**k for k in range(2, 10)]
entropies = [[attention_stats(d, scaled=s)[0] for d in ds]
for s in (False, True)]
d2l.plot(ds, entropies, 'dimension d', 'entropy (nats)',
legend=['unscaled', 'scaled by 1/sqrt(d)'], xscale='log')
def attention_stats(d, num_keys=64, num_queries=256, scaled=True, seed=0):
key_q, key_k = jax.random.split(jax.random.key(seed))
q = jax.random.normal(key_q, (num_queries, d))
keys = jax.random.normal(key_k, (num_queries, num_keys, d))
scores = (keys @ q[..., None]).squeeze(-1)
if scaled:
scores = scores / math.sqrt(d)
alpha = jax.nn.softmax(scores, axis=-1)
entropy = -(alpha * jnp.log(alpha + 1e-12)).sum(-1).mean()
# Frobenius norm of the softmax Jacobian diag(alpha) - alpha alpha^T
jac = (alpha[..., None] * jnp.eye(num_keys)
- alpha[..., :, None] * alpha[..., None, :])
return float(entropy), float(jnp.sqrt((jac**2).sum((-2, -1))).mean())
ds = [2**k for k in range(2, 10)]
entropies = [[attention_stats(d, scaled=s)[0] for d in ds]
for s in (False, True)]
d2l.plot(ds, entropies, 'dimension d', 'entropy (nats)',
legend=['unscaled', 'scaled by 1/sqrt(d)'], xscale='log')
A uniform distribution over \(64\) keys has entropy \(\ln 64 \approx 4.2\) nats. The scaled scores retain an entropy of about \(3.7\) nats across two orders of magnitude in \(d\). Thus the weight distribution retains similar sharpness as the vector dimension grows. Without scaling, the entropy collapses as \(d\) grows—below \(0.2\) nats by \(d = 512\), a near-deterministic weighting that concentrates almost all its mass on a single key purely because the vectors are long. The Jacobian norm measures the corresponding change in the score gradient:
torch.manual_seed(0)
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.259 (unscaled), 0.178 (scaled)
d = 64: Jacobian norm 0.208 (unscaled), 0.182 (scaled)
d = 512: Jacobian norm 0.093 (unscaled), 0.184 (scaled)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)With scaling, the Jacobian norm is the same at every dimension. Without it, the norm decays as the softmax saturates—by \(d = 512\) it has fallen to about half the scaled value and is still shrinking. At realistic dimensions, a model with unscaled scores would therefore begin with near-one-hot attention and small score gradients. Division by \(\sqrt{d}\) removes the dimension-induced scaling, so it is part of the definition Equation 10.2.1 rather than a tuned hyperparameter.
10.2.2 Masking
Attention as defined so far attends to every key. In practice we routinely need it not to, for one of two reasons. The first is padding: sequences of different lengths end up in the same minibatch, padded with dummy tokens to a common length (here shown as <blank>):
Dive into Deep Learning
Learn to code <blank>
Hello world <blank> <blank>Padding tokens carry no meaning, and no query should waste weight on them. The second reason is causality: a language model trained to predict the next token computes outputs for all positions of a sequence in parallel, and the query at position \(t\) must not see keys at positions beyond \(t\)— otherwise the model can copy the very future it is being trained to predict. Both cases call for the same operation: restrict the attention sum \(\sum_{i=1}^n \alpha(\mathbf{q}, \mathbf{k}_i) \mathbf{v}_i\) to a valid prefix \(\sum_{i=1}^l \alpha(\mathbf{q}, \mathbf{k}_i) \mathbf{v}_i\) with \(l \leq n\), where \(l\) depends on the sequence (padding) or on the query position (causality).
10.2.2.1 The Masked Softmax Operation
The masked softmax implements this restriction without branching over keys. It overwrites invalid scores with a very negative number before the softmax, making their normalized weights zero. This formulation preserves the regular array operations used by optimized batched kernels.
The replacement value depends on the numerical precision. Older codebases used a literal constant such as \(-10^{6}\). In single precision that works; in the half precisions that modern training runs in, it does not. The float16 format tops out near \(6.5 \times 10^4\), so \(-10^6\) silently overflows, and in bfloat16 a merely-large constant may fail to fully suppress a weight once genuine scores are large themselves. Writing literal \(-\infty\) masks exactly, but if every key of some query is masked, the softmax returns NaN and can invalidate the training run. We therefore adopt the dtype-safe idiom, which masks with the most negative finite value of the score’s dtype (torch.finfo(X.dtype).min and jnp.finfo(X.dtype).min, respectively): the masked weights are exactly zero at any precision, with no NaN. A fully masked query, which has no valid key, would otherwise produce a uniform average over invalid values. The masked_softmax implementation therefore zeroes such a row, although callers should still ensure that every query retains at least one valid key.
def masked_softmax(X, valid_lens):
"""Perform softmax operation by masking elements on the last axis."""
# X: 3D tensor, valid_lens: 1D or 2D tensor
if valid_lens is None:
return F.softmax(X, dim=-1)
shape = X.shape
if valid_lens.dim() == 1:
valid_lens = torch.repeat_interleave(valid_lens, shape[1])
else:
valid_lens = valid_lens.reshape(-1)
mask = torch.arange(shape[-1], device=X.device)[None, :]
mask = mask < valid_lens[:, None]
# Most negative finite score: exactly zero weight after the softmax,
# at any precision, without the NaN risk of literal -inf
X = X.reshape(-1, shape[-1]).masked_fill(~mask, torch.finfo(X.dtype).min)
weights = F.softmax(X, dim=-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 = torch.where(mask.any(-1, keepdim=True), weights, 0.0)
return weights.reshape(shape)def masked_softmax(X, valid_lens):
"""Perform softmax operation by masking elements on the last axis."""
# X: 3D tensor, valid_lens: 1D or 2D tensor
if valid_lens is None:
return jax.nn.softmax(X, axis=-1)
shape = X.shape
if 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)To illustrate how this function works, consider a minibatch of two examples with two queries and four keys each, where the valid lengths are \(2\) and \(3\), respectively. All weights beyond the valid length come out as zero, and each row still sums to \(1\):
masked_softmax(torch.rand(2, 2, 4), torch.tensor([2, 3]))tensor([[[0.7005, 0.2995, 0.0000, 0.0000],
[0.5080, 0.4920, 0.0000, 0.0000]],
[[0.2296, 0.3997, 0.3707, 0.0000],
[0.1914, 0.4130, 0.3956, 0.0000]]])masked_softmax(jax.random.uniform(jax.random.key(0), (2, 2, 4)),
jnp.array([2, 3]))Array([[[0.4922724 , 0.5077276 , 0. , 0. ],
[0.5997409 , 0.40025908, 0. , 0. ]],
[[0.3531402 , 0.19856434, 0.4482954 , 0. ],
[0.31042197, 0.2873157 , 0.40226236, 0. ]]], dtype=float32)For finer control we can pass a two-dimensional tensor of valid lengths, one per query:
masked_softmax(torch.rand(2, 2, 4), torch.tensor([[1, 3], [2, 4]]))tensor([[[1.0000, 0.0000, 0.0000, 0.0000],
[0.3538, 0.3878, 0.2584, 0.0000]],
[[0.5266, 0.4734, 0.0000, 0.0000],
[0.2870, 0.2921, 0.2345, 0.1864]]])masked_softmax(jax.random.uniform(jax.random.key(1), (2, 2, 4)),
jnp.array([[1, 3], [2, 4]]))Array([[[1. , 0. , 0. , 0. ],
[0.35672572, 0.2543091 , 0.38896522, 0. ]],
[[0.43679553, 0.5632045 , 0. , 0. ],
[0.2580636 , 0.2509437 , 0.258296 , 0.23269679]]], dtype=float32)10.2.2.2 Causal Masking
Per-query valid lengths are exactly what causality needs: for a sequence of length \(n\), the query at position \(t\) may attend to keys \(1, \ldots, t\), so the valid lengths are \((1, 2, \ldots, n)\), shared by every sequence in the batch. The resulting attention pattern is lower triangular:
torch.manual_seed(0)
scores = torch.randn(1, 6, 6)
causal_lens = torch.arange(1, 7)[None, :] # query t sees keys 1..t
d2l.show_heatmaps(masked_softmax(scores, causal_lens)[None],
xlabel='Keys', ylabel='Queries')
scores = jax.random.normal(jax.random.key(2), (1, 6, 6))
causal_lens = jnp.arange(1, 7)[None, :] # query t sees keys 1..t
d2l.show_heatmaps(masked_softmax(scores, causal_lens)[None],
xlabel='Keys', ylabel='Queries')
This triangular mask lets a generative model train all sequence positions in parallel without exposing future tokens. Generation additionally requires a shifted next-token objective and a decoding loop. The mask will accompany us through every decoder in the chapters ahead.
10.2.2.3 Composing Masks
valid_lens describes prefixes, which cover the two cases above, but the general interface is a boolean tensor: entry \((i, j)\) says whether query \(i\) may attend to key \(j\). Every requirement takes this form — padding excludes keys beyond the sequence length, causality excludes keys beyond the query, and structural patterns such as the attention windows of Section 10.5 exclude by distance — and a key must survive all requirements at once, so masks compose by logical AND. Broadcasting keeps the bookkeeping cheap: a padding mask has shape \((\textrm{batch}, 1, \textrm{keys})\), a causal mask \((1, \textrm{queries}, \textrm{keys})\), and their AND broadcasts to the full \((\textrm{batch}, \textrm{queries}, \textrm{keys})\) without materializing either input per example. Applying the composite is the same idiom as before: overwrite the excluded scores with the dtype’s most negative finite value, then softmax.
valid_lens, n = torch.tensor([6, 3]), 6
j = torch.arange(n)
padding = (j[None, :] < valid_lens[:, None])[:, None, :] # (batch, 1, key)
causal = (j[None, :] <= j[:, None])[None, :, :] # (1, query, key)
mask = padding & causal # (batch, query, key)
torch.manual_seed(0)
scores = torch.randn(2, n, n)
weights = F.softmax(
scores.masked_fill(~mask, torch.finfo(scores.dtype).min), dim=-1)
d2l.show_heatmaps(weights[:, None], xlabel='Keys', ylabel='Queries')
valid_lens, n = jnp.array([6, 3]), 6
j = jnp.arange(n)
padding = (j[None, :] < valid_lens[:, None])[:, None, :] # (batch, 1, key)
causal = (j[None, :] <= j[:, None])[None, :, :] # (1, query, key)
mask = padding & causal # (batch, query, key)
scores = jax.random.normal(jax.random.key(0), (2, n, n))
weights = jax.nn.softmax(
jnp.where(mask, scores, jnp.finfo(scores.dtype).min), axis=-1)
d2l.show_heatmaps(weights[:, None], xlabel='Keys', ylabel='Queries')
The first sequence shows the plain causal triangle; the second is cut off at its valid length of \(3\), the intersection of both constraints. Composition sharpens the fully-masked hazard flagged above: masks that are harmless alone can leave some query with an empty intersection, so the guarantee of at least one valid key per query must hold for the composite. The same mask composition handles packed sequences, in which several documents are concatenated into one training row. ANDing the causal mask with a block-diagonal mask prevents attention across document boundaries.
10.2.3 Batched Attention
10.2.3.1 Batch Matrix Multiplication
Attention is computed on minibatches of queries, keys, and values, so we need to multiply batches of matrices by one another. Assume that
\[ \mathbf{Q} = [\mathbf{Q}_1, \mathbf{Q}_2, \ldots, \mathbf{Q}_n] \in \mathbb{R}^{n \times a \times b}, \qquad \mathbf{K} = [\mathbf{K}_1, \mathbf{K}_2, \ldots, \mathbf{K}_n] \in \mathbb{R}^{n \times b \times c}. \]
Then the batch matrix multiplication (BMM) computes one matrix product per batch element,
\[\textrm{BMM}(\mathbf{Q}, \mathbf{K}) = [\mathbf{Q}_1 \mathbf{K}_1, \mathbf{Q}_2 \mathbf{K}_2, \ldots, \mathbf{Q}_n \mathbf{K}_n] \in \mathbb{R}^{n \times a \times c}. \tag{10.2.3}\]
A framework batch-matrix multiplication has the expected shape:
Q = torch.ones((2, 3, 4))
K = torch.ones((2, 4, 6))
d2l.check_shape(torch.bmm(Q, K), (2, 3, 6))Q = jnp.ones((2, 3, 4))
K = jnp.ones((2, 4, 6))
d2l.check_shape(jax.lax.batch_matmul(Q, K), (2, 3, 6))10.2.3.2 The DotProductAttention Class
Scaled dot-product attention can now be written in its matrix form. For \(n\) queries and \(m\) key–value pairs, with queries and keys of length \(d\) and values of length \(v\), stack the queries into \(\mathbf{Q} \in \mathbb{R}^{n \times d}\), the keys into \(\mathbf{K} \in \mathbb{R}^{m \times d}\), and the values into \(\mathbf{V} \in \mathbb{R}^{m \times v}\). Then the entire attention computation is two matrix products and a softmax:
\[ \mathrm{softmax}\left(\frac{\mathbf{Q} \mathbf{K}^\top }{\sqrt{d}}\right) \mathbf{V} \in \mathbb{R}^{n \times v}. \tag{10.2.4}\]
Requiring queries and keys to share the length \(d\) is no real restriction: a learned matrix \(\mathbf{M}\) turns \(\mathbf{q}^\top \mathbf{k}\) into \(\mathbf{q}^\top \mathbf{M} \mathbf{k}\) and translates between spaces of different dimension (an exercise below). Applied to a minibatch, Equation 10.2.4 uses the batch matrix multiplication of Equation 10.2.3 twice. The implementation applies dropout to the attention weights as regularization, and stores the weights for visualization:
class DotProductAttention(nn.Module):
"""Scaled dot product attention."""
def __init__(self, dropout):
super().__init__()
self.dropout = nn.Dropout(dropout)
# 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 forward(self, queries, keys, values, valid_lens=None):
d = queries.shape[-1]
# Swap the last two dimensions of keys with keys.transpose(1, 2)
scores = torch.bmm(queries, keys.transpose(1, 2)) / math.sqrt(d)
self.attention_weights = masked_softmax(scores, valid_lens)
return torch.bmm(self.dropout(self.attention_weights), values)class DotProductAttention(nnx.Module):
"""Scaled dot product attention."""
def __init__(self, dropout, rngs=None):
rngs = nnx.Rngs(dropout=0) if rngs is None else rngs
self.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 self
return self.dropout(attention_weights) @ values, attention_weightsTo see the class at work we use a minibatch of size \(2\), with \(10\) keys and values of dimension \(4\), a single \(2\)-dimensional query per example, and valid lengths of \(2\) and \(6\). The output is one \(4\)-dimensional row per query:
torch.manual_seed(0)
queries = torch.randn(2, 1, 2)
keys = torch.randn(2, 10, 2)
values = torch.randn(2, 10, 4)
valid_lens = torch.tensor([2, 6])
attention = DotProductAttention(dropout=0.5)
attention.eval()
d2l.check_shape(attention(queries, keys, values, valid_lens), (2, 1, 4))queries = jax.random.normal(jax.random.key(0), (2, 1, 2))
keys = jax.random.normal(jax.random.key(1), (2, 10, 2))
values = jax.random.normal(jax.random.key(2), (2, 10, 4))
valid_lens = jnp.array([2, 6])
attention = DotProductAttention(dropout=0.5)
output, attention_weights = nnx.view(
attention, deterministic=True)(queries, keys, values, valid_lens)
d2l.check_shape(output, (2, 1, 4))The stored attention weights confirm the effect of the mask: weights vanish beyond the second and sixth key, respectively.
d2l.show_heatmaps(attention.attention_weights.reshape((1, 1, 2, 10)),
xlabel='Keys', ylabel='Queries')
d2l.show_heatmaps(attention_weights.reshape((1, 1, 2, 10)),
xlabel='Keys', ylabel='Queries')
10.2.4 Attention as Learned Alignment
Learned attention first became prominent in neural machine translation. Around 2014, the leading neural approach encoded a source sentence with an RNN into a single fixed-size state vector and decoded the translation from that vector with a second RNN (we build such encoder–decoder models in full in Section 18.1). The design has the flaw this chapter opened with: one fixed-size vector must represent the entire sentence, which becomes increasingly restrictive as sentences grow. Translation quality degraded visibly with sentence length. Graves (2013) had faced a version of this problem when generating handwriting from text, and solved it with a differentiable model that aligned each output pen stroke with a position in the source text—though with the constraint that the alignment could only move forward, an assumption borrowed from decoding in speech recognition (Rabiner and Juang 1993).
Bahdanau et al. (2015) removed the constraint. Their translation model kept the two RNNs but changed the decoder input at every step. The current decoder state served as a query against all encoder states, which serve as keys and values, and feed the resulting weighted summary—a fresh one per output token—into the next prediction. The paper’s title called the idea “jointly learning to align and translate”, and the learned weights behaved exactly like the soft alignments of classical statistical translation: mostly monotone along the diagonal, with clean departures where the two languages order words differently, as sketched in Figure 10.2.1. The training objective did not directly supervise these alignments; they emerged while the model learned translation. This computation instantiates the attention mechanism in Equation 10.1.1.
One detail differed from the scoring function we settled on above. The decoder state and the encoder states were vectors of different sizes, so instead of a dot product, Bahdanau et al. (2015) scored with a small one-hidden-layer MLP, now known as additive attention:
\[a(\mathbf{q}, \mathbf{k}) = \mathbf{w}_v^\top \tanh(\mathbf{W}_q \mathbf{q} + \mathbf{W}_k \mathbf{k}) \in \mathbb{R}, \tag{10.2.5}\]
where \(\mathbf{W}_q \in \mathbb{R}^{h \times q}\), \(\mathbf{W}_k \in \mathbb{R}^{h \times k}\), and \(\mathbf{w}_v \in \mathbb{R}^{h}\) are learned. The two projections embed queries and keys into a shared \(h\)-dimensional space, and \(\mathbf{w}_v\) reads a score off the sum. It takes only a few lines to compute—here scoring three \(20\)-dimensional queries against six \(2\)-dimensional keys, dimensions no dot product could pair up:
torch.manual_seed(0)
queries, keys = torch.randn(3, 20), torch.randn(6, 2)
num_hiddens = 8
W_q = torch.randn(num_hiddens, 20) / math.sqrt(20)
W_k = torch.randn(num_hiddens, 2) / math.sqrt(2)
w_v = torch.randn(num_hiddens) / math.sqrt(num_hiddens)
features = torch.tanh((queries @ W_q.T)[:, None, :]
+ (keys @ W_k.T)[None, :, :])
scores = features @ w_v
d2l.check_shape(scores, (3, 6))
F.softmax(scores, dim=-1)tensor([[0.1692, 0.1612, 0.1816, 0.1560, 0.1502, 0.1818],
[0.1664, 0.1849, 0.1657, 0.1338, 0.1623, 0.1870],
[0.1517, 0.2188, 0.1441, 0.1317, 0.2166, 0.1371]])key_q, key_k, k1, k2, k3 = jax.random.split(jax.random.key(0), 5)
queries = jax.random.normal(key_q, (3, 20))
keys = jax.random.normal(key_k, (6, 2))
num_hiddens = 8
W_q = jax.random.normal(k1, (num_hiddens, 20)) / math.sqrt(20)
W_k = jax.random.normal(k2, (num_hiddens, 2)) / math.sqrt(2)
w_v = jax.random.normal(k3, (num_hiddens,)) / math.sqrt(num_hiddens)
features = jnp.tanh((queries @ W_q.T)[:, None, :]
+ (keys @ W_k.T)[None, :, :])
scores = features @ w_v
d2l.check_shape(scores, (3, 6))
jax.nn.softmax(scores, axis=-1)Array([[0.13390937, 0.3772766 , 0.17127123, 0.1027277 , 0.06276076,
0.15205443],
[0.19218282, 0.31524566, 0.14511336, 0.08953383, 0.07662669,
0.18129766],
[0.15669951, 0.3378386 , 0.15281527, 0.11713327, 0.09428051,
0.1412329 ]], dtype=float32)Dot-product attention later became the standard choice partly because it maps efficiently to accelerator hardware. All query–key dot products form one matrix multiplication, whereas the additive score requires an \(n \times m \times h\) tensor of hidden activations. When a learned metric is wanted, we can project queries and keys with learned matrices before the dot product and have it at matmul speed—which is precisely the form attention takes inside the Transformer, whose authors then discarded the RNN scaffolding altogether and kept attention as the only mechanism relating sequence positions (Vaswani et al. 2017). The next next sections develop multiple attention heads and then introduce explicit positional information to replace the RNN’s sequential ordering.
10.2.5 Summary
Scaled dot-product attention computes scores with one batched matrix multiplication. Dividing by \(\sqrt d\) keeps the score variance independent of dimension and prevents premature softmax saturation. A mask excludes padding, future positions, or other invalid query–key pairs before normalization; use the most negative finite value of the dtype so the operation remains valid in reduced precision. Boolean masks for several constraints compose by logical AND. DotProductAttention combines scoring, masking, dropout, and value pooling. Additive attention provides an alternative learned score and was the form used in the original neural alignment model.
10.2.6 Exercises
- Implement distance-based attention by modifying the
DotProductAttentioncode. You only need the squared norms of the keys \(\|\mathbf{k}_i\|^2\) for an efficient implementation. - Modify dot-product attention to allow for queries and keys of different dimensionalities by employing a matrix \(\mathbf{M}\) to adjust dimensions, scoring with \(\mathbf{q}^\top \mathbf{M} \mathbf{k}\).
- How does the computational cost of Equation 10.2.4 scale with the dimensionality of keys, queries, and values, and with their number? What about the memory bandwidth requirements?
- Derive the softmax Jacobian \(\mathrm{diag}(\boldsymbol{\alpha}) - \boldsymbol{\alpha}\boldsymbol{\alpha}^\top\) and verify it numerically against automatic differentiation on a random score vector. Compute its Frobenius norm when \(\boldsymbol{\alpha}\) is one-hot and when it is uniform over \(m\) keys. Which regime does the saturation experiment approach as \(d\) grows without scaling?
- What does
masked_softmaxreturn for a query whose valid length is \(0\)? Compare the behavior of masking with the most negative finite value against masking with literal \(-\infty\). In what situations can a fully masked query arise in practice, and what would you do about it? - Count the parameters and floating-point operations needed to score \(n\) queries against \(m\) keys with additive attention (hidden size \(h\)) and with scaled dot-product attention (shared dimension \(d\)). Implement a batched version of the additive score and time both variants for \(d = h = 64\) and \(d = h = 256\). Which is faster on your hardware, and why?