§3.2Self-Attention: Scaled Dot-Product Attention

Self-attention is the beating heart of the transformer. This section builds it from first principles: the idea (a context-dependent weighted average), the concrete recipe (scaled dot-product attention), the PyTorch code, and an interactive lab you can poke at.

The idea: a context-dependent weighted average

Self-attention self-attention Attention where the queries, keys, and values all come from the same sequence; each output embedding is a context-dependent weighted average of the sequence. defined in ch. 3 — open in glossary lets each token’s representation depend on the whole sequence instead of being fixed. Given token embeddings $x_1, \dots, x_n$, it produces new embeddings $x_1', \dots, x_n'$ where each $x_i'$ is a weighted combination of all the inputs:

Reading it term by term:

• $x_i'$ — the updated embedding for token $i$ (same dimension as $x_i$).
• $x_j$ — the original embedding of token $j$ (one term per token in the sequence).
• $w_{ji}$ — the attention weight attention weights Normalized coefficients w_ji (each row sums to 1) that say how much each token contributes to another token's updated embedding. defined in ch. 3 — open in glossary : how much token $j$ contributes to token $i$‘s new representation.
• $\sum_j w_{ji} = 1$ — the weights over $j$ form a proper average (they’re non-negative and sum to one).

Worked example. Suppose “flies” attends with weights $w = [0.01, 0.18, 0.01, 0.01, 0.80]$ over [time, flies, like, an, arrow]. Then its updated embedding is

— mostly “arrow” with a dash of its own identity. Those weights aren’t invented; the rest of this section is about how to compute them.

Scaled dot-product attention

There are several ways to compute those weights; the standard one — from Attention Is All You Need — is scaled dot-product attention scaled dot-product attention The standard self-attention recipe: score queries against keys with a dot product, scale by 1/√d_k, softmax to get weights, then take the weighted sum of values. defined in ch. 3 — open in glossary , in four steps:

1. Project each embedding into three vectors: a query query The learned projection of a token embedding that does the 'asking' — matched against every key to decide how much to attend. defined in ch. 3 — open in glossary , a key key The learned projection of a token embedding that is 'matched against' by queries via a similarity (dot product). defined in ch. 3 — open in glossary , and a value value The learned projection of a token embedding that carries the content mixed into the output according to the attention weights. defined in ch. 3 — open in glossary .
2. Score every query against every key with a dot product (the similarity function). For $n$ tokens this yields an $n \times n$ matrix of attention scores attention scores The raw n×n dot products of queries with keys, before scaling and softmax turn them into attention weights. defined in ch. 3 — open in glossary .
3. Scale and softmax the scores into weights that sum to 1.
4. Weighted-sum the value vectors using those weights.

Step 2 — score with the dot product

The similarity between query $q_i$ and key $k_j$ is their dot product. Stacking all queries into a matrix $Q$ and all keys into $K$, every pairwise score is a single matrix multiply:

• $Q$ — the queries, shape $[n, d_k]$ (one row per token).
• $K$ — the keys, shape $[n, d_k]$; $K^{\top}$ is its transpose $[d_k, n]$.
• $S$ — the score matrix, shape $[n, n]$. $S_{ij}$ is how strongly token $i$‘s query matches token $j$‘s key.

Worked example. With $q_{\text{flies}} = [0,3,0,0]$ and $k_{\text{arrow}} = [0,3,0,0]$, the score is $0\cdot0 + 3\cdot3 + 0\cdot0 + 0\cdot0 = 9$. Against $k_{\text{time}} = [2,0,0,1]$ it is $0$. “flies” strongly prefers “arrow”.

Step 3 — scale, then softmax

Raw dot products can get large, which pushes the softmax into a region where its gradients vanish. So we divide by $\sqrt{d_k}$ first, then normalize:

• $d_k$ — the key/query dimension. Dividing by $\sqrt{d_k}$ keeps the variance of the scores around 1 regardless of dimension.
• $\operatorname{softmax}$ — turns a row of real scores into positive weights that sum to 1.
• $w_i$ — token $i$‘s attention distribution over all keys.

Worked example. The “flies” scores $[0, 6, 0, 0, 9]$ scaled by $\sqrt{4}=2$ become $[0, 3, 0, 0, 4.5]$. Softmax gives $w_{\text{flies}} \approx [0.01,\ 0.18,\ 0.01,\ 0.01,\ 0.80]$ — exactly the weights from the averaging example above. Without the $\div 2$ scaling, the raw $[0,6,0,0,9]$ would softmax to a spikier $\approx[0.00, 0.05, 0.00, 0.00, 0.95]$.

Step 4 — weighted sum of values

Finally, mix the value vectors with those weights. In matrix form the whole mechanism is one tidy expression:

• The softmax term is the $[n, n]$ weight matrix $W$.
• $V$ is the values, shape $[n, d_v]$.
• $W V$ is the output, shape $[n, d_v]$: row $i$ is $\sum_j w_{ji} v_j$.

Try it: the attention lab

Everything above happens live below. Pick a query token and step through the four operations. Watch flies collapse onto arrow at the softmax step.

And here is the whole $5 \times 5$ weight matrix at once — every row is one token’s attention distribution. Notice function words (“like”, “an”) spread their attention, while “flies” concentrates on “arrow”.

Implementing it in PyTorch

First tokenize the text into input ids (excluding the [CLS]/[SEP] special tokens for simplicity):

inputs = tokenizer(text, return_tensors="pt", add_special_tokens=False)
inputs.input_ids
# tensor([[ 2051, 10029,  2066,  2019,  8612]])

Then look up a dense embedding for each id with an nn.Embedding layer:

from torch import nn
from transformers import AutoConfig

config = AutoConfig.from_pretrained(model_ckpt)
token_emb = nn.Embedding(config.vocab_size, config.hidden_size)  # Embedding(30522, 768)
inputs_embeds = token_emb(inputs.input_ids)
inputs_embeds.size()
# torch.Size([1, 5, 768])   ->  [batch_size, seq_len, hidden_dim]

Now score, scale, softmax, and mix. Keeping query = key = value = inputs_embeds for simplicity:

import torch
from math import sqrt
import torch.nn.functional as F

query = key = value = inputs_embeds
dim_k = key.size(-1)                                   # 768
scores = torch.bmm(query, key.transpose(1, 2)) / sqrt(dim_k)
scores.size()                                          # torch.Size([1, 5, 5])

weights = F.softmax(scores, dim=-1)
weights.sum(dim=-1)                                    # tensor([[1., 1., 1., 1., 1.]])

attn_outputs = torch.bmm(weights, value)
attn_outputs.shape                                     # torch.Size([1, 5, 768])

That’s the entire mechanism — two matrix multiplies and a softmax. Wrapped up as a reusable function:

def scaled_dot_product_attention(query, key, value):
    dim_k = query.size(-1)
    scores = torch.bmm(query, key.transpose(1, 2)) / sqrt(dim_k)
    weights = F.softmax(scores, dim=-1)
    return torch.bmm(weights, value)

Follow the shapes end to end:

Why three separate vectors?

With query = key = value set equal, a token’s query matches itself perfectly — the dot product of a vector with itself is large — so every token attends mostly to its own position. But a word’s meaning is usually better informed by its neighbours than by a second copy of itself: “flies” is clarified by “time” and “arrow”, not by “flies” again.

The fix is to give the layer freedom: apply three different learned linear projections to each embedding, producing distinct query, key, and value vectors that can specialize. Stacking several such projections in parallel gives multi-headed attention — the subject of the next section.