← All tutorials

Attention

How each token decides which tokens to attend to

On this page

Attention computes a new vector for every token as a weighted average of all the tokens' vectors. The weights say how much each token attends to, or uses, each other token. Training learns how to set those weights. Several copies of attention, called heads, run side by side in one layer, each with its own learned weights. Attention runs in every layer on every forward pass, over all the tokens at once in training.

1 Long dependencies break sequential models

Older sequence models like RNNs (recurrent neural networks) read one token at a time, a token being a word or part of a word, so they cannot run in parallel. If the word that decides the meaning showed up 50 words ago, an RNN carries it through 50 sequential steps. Every step is a chance to forget.

CNNs (convolutional neural networks, which slide a small learned filter across the input) see only a local neighborhood, so many layers must stack before one end of a sentence can reach the other.

In "the cat sat on the mat because it was tired", attention lets "it" look straight at "cat" in one step. That is self-attention.

2 Each token learns which other tokens to attend to

In "the cat sat on the mat", when the model processes "sat", which tokens should it attend to? Probably "cat" more than "the". Attention is the model learning those attention weights by itself. Each token gets a query, a key and a value; stacked over tokens they form the matrices Q, K and V. The query is what the token looks for, a key is what a token offers to be matched on, and a value is what it hands over when it is picked, like search text, page titles and the pages themselves.

All three come from the input embedding matrix X: one row per token, each row a learned vector of 512 numbers. One learned weight matrix per role turns X into queries, keys and values. Queries and keys need the same length, because each query is multiplied entry by entry with each key and summed (the dot product of Section 3). With 6 tokens and a query length of 64:

3 Scores become weights, and the weights average the values

Attention turns each token's query into one output vector in three moves: score it against every key, turn the scores into weights, and take the weighted average of the values. The score of a query against a key is their dot product: multiply the two vectors entry by entry and add up. A softmax turns each query's row of scores into weights that are positive and sum to 1. The token's output is the average of the values under those weights.

The attention equation in PyTorch, starting from queries, keys and values that are already projected: six tokens go in, one output row per token comes out, and the three lines inside the function carry the idea.

import torch

def attention(Q, K, V):
    d_k = Q.size(-1)
    S = Q @ K.transpose(-2, -1) / d_k ** 0.5  # (n, n)
    A = S.softmax(dim=-1)
    return A @ V                               # (n, d_v)

n, d_k, d_v = 6, 64, 32
Q, K = torch.randn(n, d_k), torch.randn(n, d_k)
V = torch.randn(n, d_v)
print(attention(Q, K, V).shape)  # torch.Size([6, 32])

1. 1Scores each query against each key, eq card 1 in Section 3, line 4.
2. 2Softmax over each row gives the weights, eq card 2 in Section 3, line 1.
3. 3Averages the values with those weights, eq card 2 in Section 3, line 2.
4. 4Six random tokens go in and one output row per query comes out, n rows of length d_v, the out of eq card 2 in Section 3, line 2.

Before the softmax, the scores are divided by √d_k, where d_k is the length of each query and key. A dot product adds one term per dimension, so scores grow with d_k. Big scores make the softmax nearly one-hot (one weight near 1, the rest near 0), and the gradients, the signals training uses to adjust the weights, become tiny. The paper's authors give this reasoning as a likely cause, not a proof.

If the entries of a query and a key are independent with mean 0 and variance 1 (variance is the average squared distance from the mean), their dot product has variance d_k, and dividing by √d_k brings it back to 1; the variance snippet after this paragraph checks it on 10,000 random pairs at d_k = 64. Look at the two printed variances, about d_k before the division and about 1 after.

import torch

NUM_PAIRS, d_k = 10_000, 64
q = torch.randn(NUM_PAIRS, d_k)
k = torch.randn(NUM_PAIRS, d_k)
dots = (q * k).sum(dim=-1)

print(dots.var())                 # about 64, which is d_k
print((dots / d_k ** 0.5).var())  # about 1

1. 1Dot products of 10,000 random query and key pairs, the numerator in eq card 1 in Section 3, line 4.
2. 2Their variance is about d_k, here 64, eq card 1 in Section 3, line 4.
3. 3Dividing by √d_k brings it back to about 1, the denominator in eq card 1 in Section 3, line 4.

4 One attention layer, walked through

Real models split text into subword pieces, parts of a word, with 512 numbers per token, and add position information to each vector so the model knows where the token sits; Positional Encoding and Masking explains it.

5 One head learns one relationship, so run several in parallel

One set of attention weights captures one type of relationship. But "cat" relates to "the" by grammar and to "sat" by meaning. One head can't do both well.

Multi-head attention runs several copies of attention side by side, called heads, each with its own projections into a smaller subspace, a slice of the model's dimensions. The 2017 transformer uses a model width d_model of 512 and h = 8 heads, so each head works in 64 dimensions.

Each head runs the attention line from Section 3 on its own projections:

Multi-head attention in PyTorch: one projection per role, cut into h heads, then the built-in scaled_dot_product_attention runs the attention equation of Section 3 in every head.

import torch
import torch.nn as nn
import torch.nn.functional as F

B, T, d, h = 2, 6, 512, 8  # batch, tokens, width, heads
d_k = d // h
X = torch.randn(B, T, d)
W_Q, W_K, W_V, W_O = (nn.Linear(d, d, bias=False) for _ in range(4))

def split(M):
    # (B, T, d) to (B, h, T, d_k): column block i goes to head i
    return M.view(B, T, h, d_k).transpose(1, 2)

Qh, Kh, Vh = split(W_Q(X)), split(W_K(X)), split(W_V(X))
H = F.scaled_dot_product_attention(Qh, Kh, Vh)
H = H.transpose(1, 2).reshape(B, T, d)  # concat the heads
out = W_O(H)                            # (B, T, d)

1. 1The tokens X and the four learned matrices, eq card in Section 5, lines 2 and 4.
2. 2Projects X once per role and slices the columns into h heads; slice i is XW_Q,i, eq card in Section 5, line 2.
3. 3Runs attention in every head at once, eq card in Section 5, line 1.
4. 4Joins the heads back side by side, eq card in Section 5, line 3.
5. 5Mixes them with W_O, eq card in Section 5, line 4.

Because each head scores in its own subspace, two heads can attend to different tokens. In "The animal didn't cross the street because it was too tired", "it" needs coreference, finding which noun a pronoun refers to, and the local clause around it.

More heads do not mean more parameters: with model width d (short for d_model), each head's slice d_k is d divided by the head count h.

Every head on this page takes its queries, keys and values from the same sequence, which makes it self-attention. When the queries come from one stream and the keys and values from another, say text asking about an image, it becomes cross-attention, covered in Cross-Attention.

6 One PyTorch call runs a whole attention layer

One self-attention layer on a padded batch, first as PyTorch's module, then by hand with the module's own projection weights and the fused function, so out2 matches out.

import torch
import torch.nn as nn
import torch.nn.functional as F

B, n, d_model, h = 2, 6, 512, 8
x = torch.randn(B, n, d_model)
pad = torch.tensor([[False] * 6, [False] * 4 + [True] * 2])  # True = padding

mha = nn.MultiheadAttention(d_model, h, batch_first=True)
out, _ = mha(x, x, x, key_padding_mask=pad, need_weights=False)

d_k = d_model // h
q, k, v = F.linear(x, mha.in_proj_weight, mha.in_proj_bias).chunk(3, dim=-1)
q, k, v = (t.view(B, n, h, d_k).transpose(1, 2) for t in (q, k, v))
keep = ~pad[:, None, None, :]  # True = may attend
heads = F.scaled_dot_product_attention(q, k, v, attn_mask=keep)
out2 = mha.out_proj(heads.transpose(1, 2).reshape(B, n, d_model))

1. 1A batch of two sequences, the input X of eq card 1 in Section 3, line 1; True marks the second one's 2 padding tokens.
2. 2The module runs the projections, all heads and W_O in one call, eq card in Section 5, lines 1 to 4.
3. 3The same projections by hand with the module's weights, cut into heads, eq card in Section 5, line 2.
4. 4The fused function runs attention in every head, eq card in Section 5, line 1; keep is True where attending is allowed.
5. 5Joins the heads and mixes them with the module's out_proj, eq card in Section 5, lines 3 and 4.

• batch_first defaults to False, meaning (sequence, batch, feature).
• need_weights=False is needed for the fused path; returned weights are averaged over heads.
• In key_padding_mask, True hides a key.
• In the function's boolean attn_mask, True keeps a key.
• The function ignores train and eval mode, so pass dropout_p only during training.

Attention as built here has no idea which token came first. How does it learn order, and how does a decoder avoid seeing later tokens? Positional Encoding and Masking.

7 References

1. Vaswani, A., et al. (2017). Attention Is All You Need. NeurIPS 2017. arXiv:1706.03762.