Self-Attention in Transformers: Explained Visually

Why do we need attention?

Before 2017, language models used RNNs and LSTMs. They read one word at a time and squeezed everything into a single memory vector. That caused two problems:

1. Long sentences were forgotten. Information from the beginning faded by the end.
2. Training was slow. Word 50 could not be processed until words 1 to 49 were done, so GPUs could not work in parallel.

The paper "Attention Is All You Need" (Vaswani et al., 2017) removed recurrence completely. In a Transformer, every word connects directly to every other word in one step, and all words are processed at the same time.

Self-attention in action

Click any word below to see how much attention it pays to each other word. Then switch the last word between tired and wide and click it again.

When the sentence ends in tired, it attends to animal (animals get tired). When it ends in wide, it attends to street (streets are wide). The word "it" has no fixed meaning. Attention builds its meaning from context, and this is why Transformers understand language so well.

Query, Key and Value: the core idea

For every word, self-attention creates three vectors by multiplying the word's embedding with three learned weight matrices:

The attention formula

All of self-attention fits in one line:

Read it from the inside out:

1. Q · K^T: compare every query with every key using a dot product. A large dot product means a good match. For n words, this gives an n × n table of scores.
2. ÷ √d_k: scale the scores down (d_k is the size of each key vector).
3. softmax: turn each row of scores into percentages that add up to 1. These are the attention weights.
4. · V: each word's output is the weighted average of all value vectors.

Step through the exact calculation for a three-word sentence. Every number below is computed live:

Why divide by √dₖ?

A dot product adds up d_k multiplications. If each number in Q and K has a variance of about 1, the dot product has a variance of about d_k. With d_k = 64, scores can easily be ±20.

Softmax of large numbers is extreme: one word gets almost 100% and the rest get almost 0%. In that region the softmax curve is flat, so its gradient is nearly zero and the model stops learning. Dividing by √d_k (8 for d_k = 64) brings the variance back to about 1 and keeps the softmax soft.

Self-attention in NumPy

This is the same three-word example as the step-through visual above.

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)

words = ["I", "love", "AI"]
X = np.array([[1, 0, 1, 0],      # "I"
              [0, 2, 0, 1],      # "love"
              [1, 1, 1, 1]])     # "AI"

W_q = np.array([[1, 0], [0, 1], [1, 0], [0, 1]])
W_k = np.array([[0, 1], [1, 0], [0, 1], [1, 0]])
W_v = np.array([[1, 0], [0, 1], [1, 1], [0, 0]])

Q, K, V = X @ W_q, X @ W_k, X @ W_v
d_k = K.shape[-1]
scores = Q @ K.T / np.sqrt(d_k)          # how well each query matches each key
weights = softmax(scores)                # each row sums to 1
output = weights @ V                     # weighted mix of values

print("Attention weights:\n", weights.round(2))
print("Output:\n", output.round(2))

Attention weights:
 [[0.01 0.8  0.19]
 [0.5  0.01 0.5 ]
 [0.05 0.19 0.77]]
Output:
 [[0.41 1.99]
 [1.99 1.5 ]
 [1.63 1.95]]

Read the first row: the word I puts 80% of its attention on love and 19% on AI. Its new vector [0.41, 1.99] is mostly made of the value of "love". In a real model, the weight matrices are learned during training, so these patterns become meaningful.

Multi-head attention

One attention calculation can only focus on one kind of relationship at a time. Language has many: grammar (subject and verb), meaning (it and animal), position (neighbouring words) and more. So Transformers run several attention heads in parallel, each with its own W_Q, W_K and W_V:

In the original Transformer, d_model = 512 and there are h = 8 heads, so each head works with d_k = 512 / 8 = 64 numbers. The total cost is about the same as one big head, but the model gets 8 different points of view.

import torch
import torch.nn as nn

mha = nn.MultiheadAttention(embed_dim=512, num_heads=8, batch_first=True)
x = torch.randn(2, 10, 512)                  # 2 sentences, 10 tokens each
out, attn = mha(x, x, x)                     # self-attention: Q, K, V all come from x

print(out.shape)
print(attn.shape)                            # averaged over heads by default
print(sum(p.numel() for p in mha.parameters()))

torch.Size([2, 10, 512])
torch.Size([2, 10, 10])
1050624

The parameter count is 4 matrices of 512 × 512 (Q, K, V and output) plus their biases: 4 × 262,144 + 4 × 512 = 1,050,624. The attention map is 10 × 10 because every token attends to every token.

Masked (causal) self-attention

A model like GPT writes text one token at a time. When it predicts word 5, it must not peek at words 6, 7 and 8. A causal mask sets the scores for all future positions to −∞ before the softmax, so their weights become exactly 0:

Encoder models such as BERT use no mask (every word sees both sides). Decoder models such as GPT, Llama and Claude use the causal mask.

Positional encoding: where is each word?

Self-attention treats a sentence like a bag of words. Without extra help, "dog bites man" and "man bites dog" produce the same set of attention scores. Transformers fix this by adding position information to each word's embedding before attention:

• Sinusoidal encoding (original Transformer): fixed sine and cosine waves of different frequencies.
• Learned position embeddings (BERT, GPT-2): one learned vector per position.
• Rotary position embedding, RoPE (Llama and most modern LLMs): rotates Q and K by an angle that depends on position, so attention scores depend on the distance between words.

Self-attention vs RNN

That last row is the main weakness of attention: doubling the context length makes the score table four times bigger. Techniques such as FlashAttention, sliding-window attention and KV caching are used to make long-context LLMs practical.