Attention

The formula

For $n$ tokens, their queries and keys (with $d_k$ numbers) and their values (with $d_v$) are stacked into matrices $Q,K\in\mathbb{R}^{n\times d_k}$ and $V\in\mathbb{R}^{n\times d_v}$:

Row by row: the score of token $i$ on token $j$ is $s_{ij}=q_i\cdot k_j/\sqrt{d_k}$; the weight is $\alpha_{ij}=e^{s_{ij}}/\sum_l e^{s_{il}}$ (each row adds up to 1) and the output is $o_i=\sum_j\alpha_{ij}\,v_j$. This is the “scaled dot-product attention” of Vaswani et al. (2017), who describe attention as a function that maps a query and a set of key-value pairs to an output: a weighted sum of the values.

In the scene, the query of “bank” is $q=(1,1)$ and the key of “river” is $k=(1,0.8)$: $q\cdot k=1\cdot1+1\cdot0.8=1.8$. The scores $1.8,\,0.7,\,0,\,-1.6$ give $e^{1.8}\approx6.0$, $e^{0.7}\approx2.0$, $e^{0}=1$, and $e^{-1.6}\approx0.2$; the sum is $\approx10$, and the weights are $60\%$, $20\%$, $10\%$, and $2\%$. The scene does not divide by $\sqrt{2}$ so the arithmetic can be followed in your head; with the division, the weights would be $46\%$, $21\%$, $13\%$, and $4\%$, and the new vector $(0.63,\,0.06)$ instead of $(0.75,\,0.05)$: the same side, less pronounced. All the numbers in the example are written by hand (see data.ts).

Where Q, K, and V come from

If $X$ is the matrix of input vectors (embeddings; in a transformer, plus position), $Q=XW^Q$, $K=XW^K$, and $V=XW^V$. Those three matrices are the parameters that are learned; the weights $\alpha_{ij}$ are not: they are computed anew for each sentence. In self-attention, $Q$, $K$, and $V$ all come from the same $X$. In the attention of Bahdanau et al. (2014), built for translation, the queries came from the decoder and the keys and values from the encoder, and the score was a small network, $e_{ij}=a(s_{i-1},h_j)$, rather than a dot product; the weights were the same softmax and the result the same weighted sum $c_i=\sum_j\alpha_{ij}h_j$. In the authors’ words, the decoder decides which parts of the source sentence to pay attention to.

Why $\sqrt{d_k}$

If the components of $q$ and $k$ are independent with mean 0 and variance 1, the product $q\cdot k=\sum_{i=1}^{d_k}q_ik_i$ has mean 0 and variance $d_k$ (footnote 4 of Vaswani et al., 2017): its typical size is $\sqrt{d_k}$, eight when $d_k=64$. Dividing by $\sqrt{d_k}$ brings the variance back to 1. Without it, softmax saturates: if $\alpha=\text{softmax}(s)$, then $\partial\alpha_i/\partial s_j=\alpha_i(\delta_{ij}-\alpha_j)$, which goes to 0 as one weight approaches 1 and the rest approach 0. That is the “slope that nearly vanishes” in the narration. With eight random keys and $d_k=64$, the largest weight averages $0.89$ unscaled and $0.36$ scaled (100 draws; the scene’s draw gives $0.98$ and $0.35$).

Several heads

According to Vaswani et al. (2017), several heads let the model attend at once to information from different representation “subspaces”; with a single head, averaging inhibits this. In their base model, $h=8$ and $d_k=d_v=512/8=64$, so eight heads cost about the same as one big one. BERT base has 12 layers of 12 heads (144) with 768 dimensions, that is, $768/12=64$ per head (Devlin et al., 2018).

The lesson’s real maps

They were computed with real-attention.py (in the lesson folder; it runs on a CPU): bert-base-uncased for English and bert-base-multilingual-cased for Spanish, both Apache-2.0 with a pinned revision. They are the weights of a single head after softmax, not averaged; the pieces of a word are added up. The “bank” head was chosen like this: among the heads in which the largest weight of “bank” falls on the word that decides its meaning in both sentences (river and money), the one with the largest sum of those two weights. Only 6 of 144 heads meet that condition in English (layer 3, head 11: 93% and 60%) and 4 of 144 in Spanish (layer 1, head 4: 44% and 65%). By contrast, the largest weight of “bank” falls on a mark ([CLS], [SEP], or the period) in 113 of 144 heads in English and 58 of 144 in Spanish. The other three are the heads with the largest average weight on the previous word, on the next word, and on the marks. Clark et al. (2019) describe these same kinds of heads in BERT: heads that attend to delimiter tokens, to a fixed positional offset, and broadly over the whole sentence.

Cost

The weight matrix is $n\times n$: the compute per layer grows as $O(n^2\cdot d)$. A recurrent layer costs $O(n\cdot d^2)$, but it needs $O(n)$ steps one after another, and information between two positions travels a path of length $O(n)$; in self-attention both are $O(1)$ (Table 1 of Vaswani et al., 2017). That is why attention makes good use of a GPU, and why long contexts are expensive.

In code

import numpy as np

def softmax(x):
    e = np.exp(x - x.max(axis=-1, keepdims=True))
    return e / e.sum(axis=-1, keepdims=True)

def attention(Q, K, V, scale=True):
    d = K.shape[-1]
    scores = Q @ K.T                     # dot products: how well each query matches each key
    if scale:
        scores = scores / np.sqrt(d)     # the piece the scene leaves for the end
    weights = softmax(scores)            # each row adds up to 1
    return weights @ V, weights          # weighted blend of the values

# "I sat on the bank of the river": the scene's numbers (d = 2)
K = np.array([[-1.6, 0], [-0.3, 1], [-0.8, -0.8], [-0.8, -0.8],
              [0.5, -0.5], [-0.8, -0.8], [-0.8, -0.8], [1, 0.8]])
V = np.array([[0, 0], [0.5, 0], [0, 0], [0, 0],
              [0.5, 0.5], [0, 0], [0, 0], [1, 0]])     # (riverbank, money)
q = np.array([[1.0, 1.0]])                          # the query of "bank"

out, w = attention(q, K, V, scale=False)
print(np.round(w, 2), np.round(out, 2))
out, w = attention(q, K, V)                          # with the division by √2
print(np.round(w, 2), np.round(out, 2))
# [[0.02 0.2  0.02 0.02 0.1  0.02 0.02 0.6 ]] [[0.75 0.05]]
# [[0.04 0.21 0.04 0.04 0.13 0.04 0.04 0.46]] [[0.63 0.06]]

And the check of $\sqrt{d}$: the standard deviation of $q\cdot k$ grows as $\sqrt{d}$; after dividing, it stays at 1.

rng = np.random.default_rng(0)
for d in (1, 2, 64, 512):
    q = rng.standard_normal((20000, d))
    k = rng.standard_normal((20000, d))
    s = (q * k).sum(axis=1)                         # 20,000 dot products q·k
    print(d, round(s.std(), 1), round((s / np.sqrt(d)).std(), 2))
# 1 1.0 1.0
# 2 1.4 1.0
# 64 8.0 1.0
# 512 22.6 1.0

What comes next

What attention alone does not do: it does not know word order (shuffle the words and each token spreads the same weights over the same tokens), and when writing text each token may only look back ($-\infty$ is placed on the scores that look at the future, before the softmax). With a position encoding, that mask, a network per token, residual connections, and normalization, blocks are stacked: the transformer, in lesson 22.