The one operation every frontier model is built around, derived from a soft dictionary lookup, with every shape, gradient, mask, and FLOP accounted for.

Chapter 3 ended with one operation in three steps. A query scores every stored vector, a softmax turns the scores into weights, and the output is the weighted average of the stored vectors (§3.6). In its barest form, the one you wrote in Exercise 3.6d, every row of $\mathbf{X}$ plays every part at once,

You also met the knobs. The score can be additive (Bahdanau) or a dot product (Luong), dividing the scores by a temperature slides the result from a plain average to a hard pick of one vector, and padding is hidden by setting its scores to $-\infty$ before the softmax.

That operation is not yet a layer that could replace an RNN, and Chapter 3 closed with a list of what it lacks. This chapter removes the items one at a time. This section takes the first two. In Exercise 3.6d (and in Bahdanau's model, where the annotations $\mathbf{h}_j$ were both scored and averaged) the same vector was what got matched and what got read out. What changes when those are two different things, as they are in every dictionary? And Chapter 3 said the soft version of the lookup is the one a network can learn. Why exactly, how strong is the learning signal, and how does temperature change it? The answers frame the rest of the chapter, from the learned projections of §4.2 to the scaling of §4.3.

Keys are what you match, values are what you get

Every program you have written uses lookups. rent["two-bedroom"] finds the key that matches and hands back its value, and the key and the value are different objects, often of different types. Let's make that numeric. We have $n$ entries. Entry $j$ has a key $\mathbf{k}_j \in \R^{d_k}$, which is what a query is compared against, and a value $\mathbf{v}_j \in \R^{d_v}$, which is what comes back when the entry is chosen. A query $\mathbf{q} \in \R^{d_k}$ arrives. A hard lookup scores every key with some similarity function, picks the winner, and returns its value,

The query and the keys must live in the same space, since they are compared. The values need not, which is why $d_v$ can differ from $d_k$. Chapter 3's attention is the special case $\mathbf{k}_j = \mathbf{v}_j$.

It helps to write the same lookup as a weighted sum. Let $\mathbf{a} \in \R^n$ be the one-hot vector $\mathbf{e}_{j^\star}$. Then

In plain words, a hard lookup is a weighted sum of all the values in which one weight is 1 and the rest are 0.

Here is a tiny example we will use throughout the section. The keys are apartment sizes, 40, 60, 80, and 100 square meters, and the values are monthly rents of 900, 1200, 1500, and 2100. The similarity is closeness, $\operatorname{sim}(q, k_j) = -(q - k_j)^2$. The question is about size and the answer is about money, so neither could stand in for the other. That is the separation of keys from values in its simplest form, and §4.2 shows why a Transformer wants it even when both come from the same token. Ask for a 69.9 m² apartment and the nearest key is 60, so you get 1200. Ask for 70.1 m² and you get 1500.

A lookup you cannot learn

A change of 0.2 m² in the query moved the answer by 300, while a small change almost anywhere else moves it by nothing at all. That staircase is the problem. The output is piecewise constant in $\mathbf{q}$ and in every $\mathbf{k}_j$, so its derivative is zero almost everywhere and undefined at the jumps. Gradient descent asks "if the query were slightly different, how would the output change?" and the hard lookup answers "not at all" until, suddenly, it changes completely. A network built on hard lookups gets no signal about what it should be looking for. There is a second, quieter problem. A hard lookup cannot blend. If the right answer lies between two entries, it cannot say so.

Soft weights, and why softmax makes them

The fix is to let the weights be any probability vector, not just a one-hot one,

In plain words, the output is a weighted average of the values, and the weights say how much each entry counts.

A weighted average with nonnegative weights that sum to 1 is a convex combination. Geometrically, $\mathbf{o}$ always lies inside the convex hull of the values, the smallest convex shape containing all the $\mathbf{v}_j$. A hard lookup sits at a corner of that shape, and a uniform average sits near its middle. Everything in between is now available.

Chapter 3 made the weights with a softmax of the scores $s_j = \operatorname{sim}(\mathbf{q}, \mathbf{k}_j)$, $a_j = e^{s_j}/\sum_{l} e^{s_l}$, and you have used the same function since §0.4. It is worth asking why this function, of all the ways to turn scores into weights. We want weights that are positive for any real score (including negative ones), that grow with the score, that are smooth so that gradients flow, and that can come as close to the hard argmax as we like. The exponential gives all four, through one property,

Every extra unit of score multiplies an entry's weight relative to another by $e \approx 2.72$. Only differences between scores matter. Adding the same constant $c$ to every score changes nothing, because

This shift invariance is more than a curiosity. It is why we can subtract the largest score before exponentiating (the stable softmax of §0.4), and in §4.2 it will tell us that one of the biases people put in attention layers does nothing at all.

Why not just divide each score by the sum of the scores? Scores can be negative, so the "weights" could be negative. The sum can be zero, so you could divide by zero. And $s_j / \sum_l s_l$ is not invariant to shifts, so the answer would depend on an arbitrary choice of zero point. The name also deserves a correction. §0.4 showed that log-sum-exp is the smooth version of the max, $\max_l s_l \le \log\sum_l e^{s_l} \le \max_l s_l + \log n$, and that its gradient is the softmax. The gradient of the max itself is the one-hot indicator of the winner, so softmax is really a soft argmax, a smooth stand-in for "which one is biggest", written as a probability vector. That is exactly what a soft lookup needs.

Back to the rents. With the similarity $s_j = -(q - k_j)^2 / (2\sigma^2)$ and a query of 75 m², the softmax weights depend on the width $\sigma$, which sets how quickly similarity falls off with distance.

With a narrow $\sigma$ the soft lookup is almost the hard one. With a wide $\sigma$ it blends the neighbors. Every entry is a smooth function of the query, so the network can learn.

Temperature decides how hard the lookup is

The width $\sigma$ in the rent example is a special case of the knob you met in §3.6. Divide every score by a temperature $\tau > 0$ before the softmax,

In plain words, a low temperature exaggerates the differences between scores, and a high temperature shrinks them.

Chapter 3's widget showed the two limits. Here is why they hold. The ratio of two weights is now $a_i/a_j = e^{(s_i - s_j)/\tau}$. Suppose the largest score $s_{j^\star}$ is unique. For any other entry, $s_i - s_{j^\star} < 0$, so as $\tau \to 0$ the ratio $a_i / a_{j^\star} = e^{(s_i - s_{j^\star})/\tau} \to 0$. All the weight flows to the winner, and the soft lookup becomes the hard one. As $\tau \to \infty$ every ratio tends to $e^0 = 1$, so the weights become uniform and the lookup returns the plain average of all the values.

What happens in between is less obvious. Here are the numbers for three scores $\mathbf{s} = (2, 1, 0.1)$, together with the entropy $H(\mathbf{a}) = -\sum_j a_j \log a_j$ of the weights, which measures how spread out they are (0 for a one-hot vector, $\log 3 \approx 1.099$ for a uniform one, see §0.5).

The entropy rises steadily with temperature. That is not a coincidence of this example. It is a theorem, and the proof is short enough to be worth knowing.

Monotonicity is what makes the temperature a well-behaved control. For any target sharpness between "one-hot" and "uniform" there is exactly one temperature that achieves it, and you can find it by bisection (Exercise 4.1d does).

You will meet temperature all over the book. Dividing attention scores by $\sqrt{d_k}$ is a fixed temperature chosen for a statistical reason, the subject of §4.3. Sampling from a language model at temperature $\tau$ is the same formula applied to the next-token logits (§6.4). And contrastive image-text models learn their temperature as a parameter, because the right sharpness for matching images to captions is not known in advance. CLIP initializes it at $\tau = 0.07$ and clips it so that the logits are never multiplied by more than 100 (Radford et al. 2021).

Drag the query in the widget below. Switch between the two similarity functions, and sweep the temperature from nearly hard to nearly uniform.

That last point deserves a sentence of its own. With dot-product scores, the length of the query acts as an inverse temperature. A network that can change the norms of its queries and keys can therefore choose, token by token, how sharp each lookup should be. This is part of why dot-product attention is so flexible, and also, as §4.3 will show, part of why it can become unstable.

Why the soft version can be learned

Now the payoff. How does the output change when one score changes? We need the derivative of a softmax weight with respect to a score, which §0.2 derived with the quotient rule as the softmax Jacobian,

where $\mathbb{1}[i=j]$ is 1 when $i = j$ and 0 otherwise. The output is $\mathbf{o} = \sum_i a_i \mathbf{v}_i$, so

In plain words, raising the score of entry $j$ pulls the output toward that entry's value, with a strength equal to the entry's current weight.

This little formula explains a lot. If entry $j$ currently has almost no weight, nudging its score barely matters. If it already has nearly all the weight, then $\mathbf{v}_j \approx \mathbf{o}$ and nudging it barely matters either. Gradients flow best when the weights are spread over a few candidates. That is our first glimpse of saturation, the failure mode the $\sqrt{d_k}$ scaling of §4.3 exists to prevent. With a temperature, the formula picks up a factor $1/\tau$, so very hot or very cold lookups both learn slowly, for opposite reasons. The hard lookup is the limit $\tau \to 0$, where the weights are one-hot and the gradient is exactly zero.

We can go one step further and differentiate with respect to the query itself, using dot-product scores $s_j = \mathbf{q}^\top\mathbf{k}_j$, so $\partial s_j / \partial \mathbf{q} = \mathbf{k}_j^\top$. By the chain rule,

where $\bar{\mathbf{k}} = \sum_j a_j \mathbf{k}_j$. The second form is allowed because $\sum_j a_j(\mathbf{v}_j - \mathbf{o}) = \mathbf{o} - \mathbf{o} = \mathbf{0}$, so subtracting $\bar{\mathbf{k}}$ adds nothing.

In plain words, the Jacobian of the output with respect to the query is the covariance between values and keys under the attention weights. The query learns to move in key directions that are associated with the values the loss wants. With $\mathbf{k}_j = \mathbf{v}_j$, as in Chapter 3, this is the covariance of the stored vectors with themselves, so a query can find an entry only by resembling what it wants to read. Separate keys let it find an entry by one property and read out another, the way the rent example finds an apartment by its size and reads its price.

Let's verify that claim numerically before believing it.

import numpy as np

def soft_lookup_batch(Q, K, V, tau=1.0):
    raise NotImplementedError

The soft lookup's gradient is not just nonzero, it points somewhere useful. The next exercise uses it to learn a query, which is the whole reason for making the lookup soft.

import numpy as np

def query_step(q, K, V, target, tau=1.0):
    z = K @ q / tau
    a = np.exp(z - z.max()); a /= a.sum()
    r = a @ V - target
    return 0.5 * r @ r, np.zeros_like(q)      # placeholder: the loss is right, the gradient is not

Look at the third test. A query that has locked onto the wrong entry at a tiny temperature has a large loss and a zero gradient. Nothing in the loss can move it. That is the hard lookup's failure, recovered as the cold limit of the soft one, and the reason trained attention keeps its scores in a moderate range.

Which similarity? Dot product, distance, or a small network

The softmax does not care where the scores come from. §3.6 compared Bahdanau's additive score with Luong's dot and general scores, and found the dot product cheapest, since a whole batch of scores is one matrix product. One more family is worth adding, because it connects attention to a much older idea, and it shows that the choice matters less than you might think.

Negative squared distance, $s_j = -\lVert\mathbf{q} - \mathbf{k}_j\rVert^2$, is what the rent example used. Expand it,

The term $\lVert\mathbf{q}\rVert^2$ is the same for every $j$, so by shift invariance it disappears inside the softmax. What is left is a dot product (times 2) plus a per-key bias $-\lVert\mathbf{k}_j\rVert^2$ that penalizes long keys. When all keys have the same norm, distance attention and dot-product attention give identical weights. This is the same identity, $\lVert\mathbf{u}-\mathbf{v}\rVert^2 = \lVert\mathbf{u}\rVert^2 + \lVert\mathbf{v}\rVert^2 - 2\mathbf{u}^\top\mathbf{v}$, that tied cosine to distance in §1.3.

Additive attention, $s_j = \mathbf{w}^\top\tanh(\mathbf{W}_1\mathbf{q} + \mathbf{W}_2\mathbf{k}_j)$, is more flexible per score but cannot be written as one matrix product. Vaswani et al. (2017) report that the two perform similarly for small $d_k$ and that additive attention wins over unscaled dot products for large $d_k$, the observation that led them to scale the dot product (§4.3).

Let's see kernel regression as attention, with three bandwidths.

A small bandwidth (low temperature) chases the noise, since each query copies its nearest neighbors. A large one (high temperature) over-smooths toward the global mean. The middle one is close to the truth. That is the bias-variance trade-off, and it is the same trade-off a network faces when it decides how sharp each attention head should be.

import numpy as np

def nadaraya_watson(xq, xs, ys, bandwidth):
    raise NotImplementedError

import numpy as np

def temperature_for_entropy(s, H_target):
    return 1.0   # placeholder: always returns tau = 1

In PyTorch

In PyTorch the soft lookup is three lines, and autograd confirms the covariance formula for the query gradient.

Read more

• Bahdanau, Cho & Bengio (2015). Neural Machine Translation by Jointly Learning to Align and Translate. The paper that introduced soft attention, as an alignment model between a decoder and an encoder.
• Radford et al. (2021). Learning Transferable Visual Models From Natural Language Supervision. CLIP, whose learned softmax temperature starts at 0.07 and is clipped at a logit scale of 100.
• Graves, Wayne & Danihelka (2014). Neural Turing Machines. Content-based addressing of a memory with a softmax over cosine similarities and a learned key strength that acts as an inverse temperature, a differentiable dictionary in its purest form.
• Zhang et al. Dive into Deep Learning, attention pooling. Nadaraya-Watson regression developed step by step into attention.
• Martins & Astudillo (2016). From Softmax to Sparsemax. A sparse alternative to softmax and its Jacobian.
• Ramsauer et al. (2020). Hopfield Networks is All You Need. Attention as the update rule of an associative memory, with temperature as inverse sharpness.

Watch next

▶

▶