Self-Attention & Transformers

See queries, keys and values combine into an attention heatmap for a sentence

AdvancedSelf-AttentionQuery/Key/ValueSoftmaxMulti-Head

Sentence (split on spaces, up to 8 tokens)

Query token (click to inspect)

Causal mask (decoder)Add positional encoding

Heads

d_model = 8, d_k = d_model / heads = 8. Embeddings are hashed from each word and the W_Q, W_K, W_V weights come from a seeded random generator, so the patterns are illustrative rather than learned. The arithmetic is exact.

Queries, keys and values

Q = X·W_Q, K = X·W_K, V = X·W_V, each 6×8. X is the embedding plus positional encoding of each token.

	0	1	2	3	4	5	6	7
the	1.32	-0.76	-1.77	0.10	-0.29	-0.23	-0.48	0.92
cat	-0.78	-1.64	-2.23	0.72	0.49	1.78	0.83	2.44
sat	-0.82	-0.22	-1.22	1.88	-0.16	0.90	0.08	1.53
on	-0.07	-0.51	-2.59	1.33	-0.58	0.91	-0.03	0.76
the	1.35	-0.06	-0.20	1.60	-0.82	-0.78	-1.68	0.08
mat	0.62	0.74	-0.56	0.22	-0.58	-0.81	-0.04	-0.52

	0	1	2	3	4	5	6	7
the	-0.78	-2.44	-1.60	-1.37	0.23	1.52	-0.63	0.37
cat	-0.44	-1.69	-1.23	-2.06	-0.26	1.27	-0.19	-1.31
sat	-0.90	-0.88	-2.30	-0.62	1.08	1.65	0.00	-0.91
on	-2.20	-0.75	-2.50	-1.83	2.18	0.42	-0.74	0.23
the	-1.90	-1.67	-2.40	-1.25	1.42	0.81	0.13	1.22
mat	-0.11	-0.07	-1.09	1.01	0.86	0.87	-0.27	0.17

	0	1	2	3	4	5	6	7
the	1.52	0.83	-1.60	1.09	0.31	-0.32	0.97	1.65
cat	0.25	-1.17	0.07	2.13	0.48	2.09	-0.12	3.29
sat	-0.85	0.48	0.09	-0.26	1.30	1.96	-1.57	0.25
on	-1.61	2.51	-0.51	0.86	0.02	1.57	-0.19	0.65
the	0.19	2.09	-1.83	-0.75	0.42	0.32	-0.12	-0.13
mat	0.15	1.48	-0.27	-0.98	0.81	-0.74	-0.80	-1.74

Scores QKᵀ / √d_k

Dot product of each query row with each key row, divided by √8 = 2.83.

scores

	the	cat	sat	on	the	mat
the	1.32	0.48	0.69	0.62	1.19	0.63
cat	3.67	1.12	2.85	3.17	4.92	1.95
sat	0.85	-0.87	0.88	0.68	1.73	1.49
on	1.82	0.58	2.06	1.37	2.25	1.64
the	-1.09	-1.46	-1.40	-2.20	-2.09	0.27
mat	-1.14	-0.52	-0.54	-0.92	-1.22	-0.20

Attention weights = softmax(scores)

Each row sums to 1. Click a row to make it the query.

What "cat" attends to

Line thickness is the attention weight. The output for this token is the weighted sum of the value rows.

z(cat) = 0.17·v(the) + 0.01·v(cat) + 0.08·v(sat) + 0.10·v(on) + 0.60·v(the) + 0.03·v(mat)

= [0.15, 1.73, -1.43, -0.19, 0.44, 0.46, -0.07, 0.28]

Transformer block

Attention mixes information between tokens; the feed-forward layer processes each token on its own.

Positional encoding

Attention ignores order, so each position adds a unique pattern of sines and cosines. Turn it off and repeated words get identical rows in Q, K and V.

dim 0 (sin, fast)dim 1 (cos, fast)dim 2 (sin, slower)dim 3 (cos, slower)Shaded columns: positions used by your sentence. Dims 4–7 change even more slowly.

How self-attention works

In self-attention every token looks at every other token in the same sentence and decides how much of each to mix into its own representation. Each token's embedding is projected three ways: a query (what am I looking for?), a key (what do I contain?) and a value (what do I pass on?).

The match between a query and a key is their dot product. Dividing by √d_k and applying softmax turns a row of scores into weights that sum to 1, and the token's output is the weighted sum of all value vectors. A Transformer stacks this with a small feed-forward network, residual connections and layer normalisation, and repeats the block many times.

The weights in this demo are random but seeded, so the patterns will not look like a trained model's (which might link "sat" strongly to "cat"). What is real is the computation: every number on screen follows the formulas below.

Key formulas

Scaled dot-product attention

Attention(Q, K, V) = softmax(QKᵀ / √d_k + M) · V

Q, K, V: X·W_Q, X·W_K, X·W_V, one row per token
d_k: length of each query and key vector
M: mask: 0 where allowed, −∞ for future tokens when causal

Softmax over one row

a_ij = exp(s_ij) / Σ_m exp(s_im)

Multi-head attention

MultiHead(X) = Concat(head₁, …, head_h) · W_O, with d_k = d_model / h

Sinusoidal positional encoding

PE(pos, 2i) = sin(pos / 10000^2i/d), PE(pos, 2i+1) = cos(pos / 10000^2i/d)

How to use this visualizer

1
Type a sentence of up to 8 words. Each word becomes a token with an 8-number embedding.
2
Click a token (or a heatmap row) to make it the query. Its row is highlighted in Q, in the scores and in the heatmap.
3
Follow the pipeline: Q, K and V, then the scaled scores, then the softmax weights, then the lines and the weighted sum that form the output.
4
Turn on the causal mask and watch the upper triangle drop to zero: each token can only attend to itself and earlier tokens.
5
Switch to 2 or 4 heads. Each head gets a smaller d_k and its own weights, so the small heatmaps show different patterns.
6
Turn positional encoding off with a repeated word (like "the" twice) and notice the two rows become identical.

Key concepts

Query, key, value

Each token asks a question (query), advertises what it contains (key) and offers content to share (value). All three are linear projections of the same embedding.

Scaling by √d_k

Dot products grow with the vector length. Dividing by √d_k keeps the scores in a range where softmax does not saturate and gradients still flow.

Causal mask

Language models must not peek at future tokens, so scores above the diagonal are set to −∞ and get zero weight after softmax.

Multi-head attention

Several smaller heads run in parallel with their own weights. One might track the previous word, another the subject of the sentence.

Residual connections

Each sub-layer adds its output to its input (Add & Norm), which keeps deep stacks of blocks trainable.

Positional encoding

Attention treats the input as a set. Adding a position-dependent signal to each embedding lets the model use word order.

Code

NumPy

import numpy as np

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

def self_attention(X, W_q, W_k, W_v, causal=False):
    Q, K, V = X @ W_q, X @ W_k, X @ W_v            # (n, d_k) each
    scores = Q @ K.T / np.sqrt(K.shape[-1])         # (n, n)
    if causal:
        n = scores.shape[0]
        scores[np.triu_indices(n, k=1)] = -np.inf    # hide future tokens
    A = softmax(scores, axis=-1)                    # rows sum to 1
    return A @ V, A                                 # outputs and weights

def positional_encoding(n, d):
    pos = np.arange(n)[:, None]
    i = np.arange(0, d, 2)[None, :]
    angle = pos / 10000 ** (i / d)
    pe = np.zeros((n, d))
    pe[:, 0::2], pe[:, 1::2] = np.sin(angle), np.cos(angle)
    return pe

Practice problems

Problem 1: Attention by hand

With d_k = 2, a query q = [1, 0] attends to two keys k₁ = [1, 1] and k₂ = [0, 2] with values v₁ = [1, 0] and v₂ = [0, 1]. Compute the scaled scores, the attention weights and the output.

Show solution

Raw scores: q·k₁ = 1, q·k₂ = 0

Scaled by √2 ≈ 1.414: s = [0.707, 0]

exp: [2.028, 1.000], sum = 3.028

weights a = [0.670, 0.330]

output = 0.670·[1, 0] + 0.330·[0, 1]

z = [0.67, 0.33]

Problem 2: Causal mask and heads

In a causal decoder the second token (index 1) of a three-token sequence has scaled scores [2, 1, 3]. What are its attention weights? Separately, a model with d_model = 512 uses 8 heads: what is d_k and what are the scores divided by?

Show solution

The third token is in the future, so its score becomes −∞: [2, 1, −∞]

exp: [7.389, 2.718, 0], sum = 10.107

weights = [0.731, 0.269, 0]

d_k = 512 / 8 = 64, so scores are divided by √64 = 8

[0.731, 0.269, 0]; d_k = 64, scale 8

Discussion

Questions, ideas or a different approach? Join the conversation.

Queries, keys and values

Scores QKᵀ / √dk

Attention weights = softmax(scores)

What "cat" attends to

Transformer block

Positional encoding

How self-attention works

Key formulas

How to use this visualizer

Key concepts

Query, key, value

Scaling by √d_k

Causal mask

Multi-head attention

Residual connections

Positional encoding

Code

Practice problems

Problem 1: Attention by hand

Problem 2: Causal mask and heads

Discussion

Scores QKᵀ / √d_k