Scaled Dot-Product Attention & Transformer Layers

Why It Matters for Modern Models

Scaled dot-product attention uses query–key scores to compute a weighted average of value vectors. Published large language models such as GPT-3 and Llama are deep stacks of decoder-only transformer blocks with causal self-attention (Brown et al., 2020; Touvron et al., 2023).
Latent diffusion (the basis of Stable Diffusion) uses a U-Net with self-attention over image latents and cross-attention from the latents to text embeddings (Rombach et al., 2022).
OpenAI describes its Sora video model as a diffusion model whose transformer operates on spacetime patches of compressed video (Brooks et al., 2024).

Questions to Think About

In what sense is each attention head a learned, input-dependent kernel over positions, and how does that differ from a fixed convolution kernel?

Interactive Demos

Row 0 (A) · unmasked projection.

One query · three sources · one continuous example

Follow one query through attention

Drag q and watch three steps happen at once: match (how far q reaches along each key), share (softmax splits the attention) and mix (the output blends the values). You can explore before answering anything; go at your own pace.

Given vectors for one unmasked head, not learned word meanings.

q = [2, 1] · Compare the query with each key.

1Match: how far does q reach along each key?

Drag anywhere in the plane, or focus q and use the arrow keys.

2Share: softmax turns scores into one pie

Score q·k / √2Share

A014%

B1.41458%

C0.70728%

Softmax exponentiates each score, then divides by the total. Only differences between scores matter.

3Mix: the output is the shares’ blend of the values

Coloured arrows add w·v for each source, tip to tail. The output can never leave the triangle. The view fits the values with equal units on both axes.

q lines up most with B’s key (score 1.414), so B takes 58% of the attention; C gets 28% and A gets 14%. Output o = [-0.004, 0.868].

Give source B more than 90% of the attention.

Watch Match → Share → Mix

A 31-second guided explanation of the starting A/B/C example. This is fixed, silent playback with captions; it does not follow or reset your experiment.

Read the same calculation

Match: query (2, 1) meets keys A = (−1, 2), B = (1, 0), C = (0, 1). Their dot products are [0, 2, 1]. Divide by √2, the square root of the key width, to get [0, 1.414, 0.707].
Share: softmax subtracts the largest score, exponentiates and divides by the sum. The weights, in A/B/C order, are [0.14, 0.576, 0.284]. They sum to one; a larger match receives more weight.
Mix: multiply values A = (2, 0), B = (0, 2), C = (−1, −1) by those weights, then add their contributions. The output is [-0.004, 0.868]. The values carry the content; the query and keys determine its weighting.

Displayed numbers are rounded. This hand-authored, unmasked example explains one attention calculation; it is not a trace from a trained model.

Show every number, the formula and the Python

Change the query · one unmasked rowQuery q, x: 2Query q, y: 1

Shared normalization: m = 1.414; Z = 1.736.

Scaled scores · s = q·k/√2

−4.504.5

Attention shares w in A, B, C order · the full strip is 1

ABCSolid dot or edge: current q. Hollow ring or dashed line: same current keys and values with q = [2, 1].

Source Aindex 0
Dot
0
Scaled
0
Share
0.14
Carry w_Av_A
[0.28, 0]
Source Bindex 1
Dot
2
Scaled
1.414
Share
0.576
Carry w_Bv_B
[0, 1.152]
Source Cindex 2
Dot
1
Scaled
0.707
Share
0.284
Carry w_Cv_C
[-0.284, -0.284]

Output o = [-0.004, 0.868]

Compared with q = [2, 1]:

Scores same; shares same; output same.

Worked example, not a test or model run. Given vectors, not learned meanings. Match the query with keys; mix values, not keys. All three sources are allowed in this view.

1 · Match and scaleq · kᵢ → sᵢ = (q · kᵢ) / √2sum(qj * kj for qj, kj in zip(q, k)) / sqrt(2)

2 · Share one denominatorwᵢ = exp(sᵢ − m) / Zexps = [exp(s - max(scores)) for s in scores]weights = [e / sum(exps) for e in exps]

3 · Carry the valuescᵢ = wᵢvᵢ; o = Σᵢ cᵢsum(w * v[j] for w, v in zip(weights, values))

m is the shared maximum scaled score; Z = exp(s_A − m) + exp(s_B − m) + exp(s_C − m). Subtracting the same maximum avoids large exponentials without changing the shares. Each carry is a weighted source value; sum the three contributions coordinate by coordinate for output o.

The reference uses the same current keys and values, with q = [2, 1]. Not every query edit changes the mixture. Solid amber borders mark numerical changes; dashed teal borders mark values unchanged within 1e-12. Labels round to three decimals, not the calculation.

Source A · index 0: k_A = [-1, 2]; v_A = [2, 0]
Source B · index 1: k_B = [1, 0]; v_B = [0, 2]
Source C · index 2: k_C = [0, 1]; v_C = [-1, -1]

Inspect precise numbers and input coordinates

Finite JavaScript numbers before display rounding; −0 is kept where supplied. Tiny allowed shares can underflow to zero; that is not a mask.

q = [2, 1]
A: key=[-1, 2]; value=[2, 0]; dot=0; score=0; share=0.14002924504337802; contribution=[0.28005849008675604, 0]
B: key=[1, 0]; value=[0, 2]; dot=2; score=1.414213562373095; share=0.575975345215362; contribution=[0, 1.151950690430724]
C: key=[0, 1]; value=[-1, -1]; dot=1; score=0.7071067811865475; share=0.28399540974126003; contribution=[-0.28399540974126003, -0.28399540974126003]
m=1.414213562373095; stable denominator=1.736185425829454
output=[-0.00393691965450399, 0.8679552806894639]

Read the same calculation in Python

This runnable Python uses the current controls and the same labels. No NumPy or random inputs are needed.

Float operands and left-to-right sums from 0.0 follow browser accumulation, not exact arithmetic. Cancellation can hide a nonzero exact change; exp/sqrt implementations and display rounding can differ across runtimes.

from math import exp, sqrt

labels = ["A","B","C"]
q = [2.0, 1.0]
keys = [[-1.0, 2.0], [1.0, 0.0], [0.0, 1.0]]
values = [[2.0, 0.0], [0.0, 2.0], [-1.0, -1.0]]

# Float operands and left-to-right sums from 0.0 follow browser accumulation.
# Cancellation can hide a nonzero exact change. This is not exact arithmetic;
# exp/sqrt implementations and display rounding can differ across runtimes.
def ordered_sum(terms):
    total = 0.0
    for term in terms:
        total += term
    return total

d_k = len(q)  # key dimension, NOT number of sources or value width
raw_scores = [ordered_sum(qj * kj for qj, kj in zip(q, k)) for k in keys]
scores = [s / sqrt(d_k) for s in raw_scores]
exps = [exp(s - max(scores)) for s in scores]
total = ordered_sum(exps)
weights = [e / total for e in exps]
output = [ordered_sum(v[j] * w for w, v in zip(weights, values))
          for j in range(len(values[0]))]

for label, raw, score, weight in zip(labels, raw_scores, scores, weights):
    print(label, round(raw, 3), round(score, 3), round(weight, 3))
print("output", [round(x, 3) for x in output])

Browser reference output: [-0.004, 0.868] (rounded to three decimals).

Inspect key and value geometry

Key space only. Each grid interval is 1 coordinate unit on both axes. Amber marks the query q you control; ink marks the three fixed keys.

Value space only. Each grid interval is 1 coordinate unit on both axes. Teal marks the mixed output o; the triangle bounds its possible positions.

Step 1 of 4 · Vectors

What does this query have in common with each key?

A, B, and C name the same three sources throughout. Compare q with their keys in key space. Their values are payloads in a different space, not arrows to compare with q.

These are optional inspection controls. The connected scores, shares and mixture stay visible in the explorable above.

Read the current numbers

Shapes: q and each k have d_k = 2 coordinates. Each v and the output have d_v = 2 coordinates. Equal widths here do not make the spaces identical.

q = [q₁, q₂]; k_A = [−1, 2]; k_B = [1, 0]; k_C = [0, 1]

i names a source; j runs over A, B, C. s is a scaled score, w its share, and o the output. Displayed decimals are rounded; calculations are not.

Keys in key space
Source	Key k	q
A	[-1, 2]	[2, 1]
B	[1, 0]	[2, 1]
C	[0, 1]	[2, 1]

Reset example restores the query and first value’s x coordinate only. All other held inputs stay unchanged.

All sources are allowed here; this is not a masked sequence output. Query edits return to row 0; value edits are shared. Other rows and the held mask stay unchanged.

Practice, help and export

A study detour counts as consultation for the attempt you leave. Clearing a response or resetting values keeps its help history.

Held example · 3 queries · 3 sources. Started from the A/B/C sequence.

Study history for this visit

These record which help you used, not completion or learning scores. Each attempt keeps its own feedback and help. Reloading the page clears this history.

Guided attention notebook · worked study · sequence-abc-v1

Core Math

If you prefer intuition first, try the demos above, then come back for the math.

Key Equation

\text{Attn}(Q,K,V) = \mathrm{softmax}\!\left(\frac{QK^\top}{\sqrt{d_k}}\right)V

Single attention head:

\text{Attn}(Q,K,V) = \mathrm{softmax}\!\left(\frac{QK^\top}{\sqrt{d_k}}\right)V

where $Q = XW_Q,\ K = XW_K,\ V = XW_V$ and $d_k$ is the number of coordinates in each query and key. If those coordinates are independent with mean 0 and variance 1, a dot product $q^\top k$ has variance $d_k$ ; dividing by $\sqrt{d_k}$ keeps the scores at unit variance so the softmax does not saturate as $d_k$ grows (Vaswani et al., 2017). Multi-head attention runs several such heads in parallel and concatenates their outputs.

A pre-norm transformer block, as in GPT-2 and many later language models (the original Transformer applied LayerNorm after each residual addition instead):

\begin{aligned} H' &= \mathrm{MHA}(\mathrm{LN}(H)) + H \\ H^{\text{out}} &= \mathrm{MLP}(\mathrm{LN}(H')) + H' \end{aligned}

Key Papers and Books

Attention Is All You Need

Vaswani et al.2017NeurIPS

Read paper →

Neural Machine Translation by Jointly Learning to Align and Translate

Bahdanau, Cho & Bengio2015ICLR

Read paper →

Understanding Deep Learning (chapter on transformers)

Prince2023MIT Press (textbook)

View book →

Connections

Where to go next

A newer notebook covers Attention, with an interactive demo.

Continue in the newer notebook

Other connections (17)

Which limitations does it address, or leave to a later idea?

Position encoding: a limitation of Attention that RoPE addresses. Open RoPE
The original Transformer added sinusoidal or learned absolute position vectors to its inputs. RoPE instead rotates queries and keys so that their dot product depends only on the offset between positions (Su et al., 2024). It does not by itself extrapolate well beyond the training length; that needed later rescaling methods such as position interpolation (Chen et al., 2023).
KV cache memory: a limitation of Attention that Efficient Attention addresses. Open Efficient Attention
During generation every attention head caches its keys and values, so the cache grows with the number of tokens, heads and layers. Multi-query and grouped-query attention shrink it by letting several query heads share one key and value head (Shazeer, 2019; Ainslie et al., 2023). FlashAttention addresses a different cost, the n × n attention matrix.
Autoregressive latency: a limitation of Attention that Speculative Decoding addresses. Open Speculative Decoding
Token-by-token decoding is the sequential bottleneck in transformer inference. Speculative decoding was invented to propose many tokens cheaply and verify them in parallel without changing the target distribution.
Production serving: a limitation of Attention that LLM Serving addresses. Open LLM Serving
Running a transformer once is easy; serving thousands of users at once within latency targets is not. Serving techniques such as continuous batching, separating prompt processing from generation, and paged key-value caches were developed to make this practical (Kwon et al., 2023).
Linear complexity: a limitation of State Space Models that Attention addresses. Open State Space Models
State-space models such as Mamba were developed as an alternative to attention's O(n²) cost: they process a sequence in O(n) time and keep a fixed-size state during generation (Gu & Dao, 2024). The price is that the whole past must be compressed into that state, which makes exact copying from long inputs harder than for transformers (Jelassi et al., 2024).
Attention costs: a limitation of Efficient Attention that Attention addresses. Open Efficient Attention
Efficient-attention methods target different costs of standard attention. FlashAttention computes exact attention without storing the n × n matrix; grouped-query attention shrinks the key-value cache; sliding-window attention limits each token to recent positions, which changes what is computed.

Where does it break down?

O(n²) memory: where Attention breaks down, taken up in Efficient Attention. Open Efficient Attention
A standard implementation stores the full n × n matrix of attention weights for the backward pass; at n = 100,000 that is 10¹⁰ numbers per head, 20 GB in 16-bit precision. FlashAttention avoids storing it by computing attention in tiles and recomputing them during the backward pass (Dao et al., 2022).
Attention dilution: where Attention breaks down, taken up in Long Context. Open Long Context
Even if memory is solved, long contexts make retrieval harder: softmax has to discriminate among many keys, and attention can become diffuse or distracted. Quality can degrade with length, not just speed.

Where else does the same idea appear?

Softmax routing: the same technique appears in MoE. Open MoE
Both use a softmax to turn scores into weights: attention over the positions in a sequence, a mixture-of-experts router over a fixed set of experts. The router is a learned linear layer, not an attention head, and it keeps only the top-scoring experts instead of averaging over all of them.
Cross-attention: the same technique appears in Diffusion. Open Diffusion
Text-to-image diffusion (Stable Diffusion) uses cross-attention to condition on text embeddings—same Q·K^T softmax mechanism, but K,V come from text while Q comes from image latents.
Dot-product logits: the same technique appears in Embeddings. Open Embeddings
Attention forms scores from query–key dot products divided by √d_k; contrastive learning forms them from embedding dot products divided by a temperature τ. Both build a matrix of dot-product similarities and apply a softmax to each row. Attention uses the resulting weights to average value vectors, while contrastive learning treats them as predicted probabilities and trains them with cross-entropy.
Softmax selection: the same technique appears in Decoding & Sampling. Open Decoding & Sampling
Attention uses softmax to select "which positions to copy from," while decoding uses softmax to select "which token to emit." Same pattern: logits → distribution → sample/argmax/weighted sum.
Cross-attention fusion: the same technique appears in Multimodal Pretraining. Open Multimodal Pretraining
Some multimodal models fuse modalities with cross-attention: in Flamingo, text tokens attend to image features through added cross-attention layers (Alayrac et al., 2022). Others use none: CLIP encodes images and text separately and only compares the two embeddings.
Sparse selection: an analogy between Attention and MoE. Open MoE
A mixture-of-experts layer sends each token to its top-k experts and ignores the rest. Softmax attention is dense, giving every position some weight, although trained heads often concentrate on a few positions; sparse-attention variants make that selection explicit.
Copy via attention: the same technique appears in Circuits. Open Circuits
An induction head completes patterns of the form [A][B] … [A] → [B]. A previous-token head in an earlier layer writes each token's predecessor into its position; the induction head then attends from the current token A to the token that followed an earlier A, and copies it (Olsson et al., 2022). Attention does the copying.
Shared attention: the same technique appears in Multimodal Pretraining. Open Multimodal Pretraining
Multimodal models reuse the attention mechanism across modalities. Some interleave image and text tokens in one sequence and use ordinary self-attention; others add cross-attention layers so that one modality attends to the other (Alayrac et al., 2022).

Which ideas are its mathematical duals?

Kernel ↔ Recurrence: State Space Models is its mathematical dual. Open State Space Models
Attention applies a (data-dependent) kernel over history; SSMs apply a structured kernel computed by recurrence/scan. Two mathematically connected ways to implement "weighted sums over the past," with different computational tradeoffs.