Attention Mechanism

Summary: The operation inside a transformer that lets token embeddings communicate with each other, so that each token’s vector can be refined based on the context it appears in.

Mechanically: for every (query, key) pair, compute a dot-product alignment score, softmax the scores into weights, and use those weights to mix value vectors into the output.

Geometrically: allow other surrounding tokens to update a given token’s word-embedding vector (of dimensionality $d_{\text{embed}}\approx 12,000$ in GPT-3), moving it from its original orientation to a different part of the high-dimensionality embedding space that better captures its semantic meaning (“fluffy” in “fluffy creature” augments the meaning of “creature” by itself)

💡 The problem attention solves

After word-embedding, every instance of “mole” has the same vector, regardless of whether it means a burrowing animal, a unit of substance, or a skin growth. Only context can disambiguate. Attention is the mechanism by which surrounding tokens inject information into each token’s vector so that downstream layers see a context-aware representation.

Crucially, this matters not just for the “interesting” tokens — it matters for the last token, because only the last vector in the sequence is used to produce the next-token prediction (see unembedding). That last vector has to absorb everything from the full context.

A single head of attention

A single “head” is parameterised by three learned matrices: $W_Q$, $W_K$, $W_V$. Given the sequence of token embeddings ${\vec{E}}_1,\dots ,{\vec{E}}_n$ flowing into the block:

This describes self-attention

Everything on this page — $Q$, $K$, $V$ all drawn from the same input token sequence — is technically the self-attention variant. GPT-style decoder-only LLMs use only self-attention.

A sibling variant, cross-attention, feeds queries from one sequence (e.g. English, or speech audio) and keys/values from another (e.g. French, or text transcript; the decoder of an encoder–decoder translation model). The machinery is otherwise identical. See self-attention-vs-cross-attention.

Step 1 — Queries and keys

Each token’s embedding, ${\vec{E}}_i\in {\mathbb{R}}^{d_{\text{embed}}}$, is projected into two small vectors, ${\vec{Q}}_i$ and ${\vec{K}}_i\in {\mathbb{R}}^{d_{\text{k}}}$. These projections are done through the query-projection matrix ($W_Q$), and the key projection matrix ($W_K$) respectively. Both matrices are of dimension ${\mathbb{R}}^{d_{\text{k}}\times d_{\text{embed}}}$ (typically $d_k\ll d_{\text{embed}}$; GPT-3 uses $d_k=128$ vs $d_{\text{embed}}=12,288$).:

${\vec{Q}}_i=W_Q{\vec{E}}_i{\vec{K}}_i=W_K{\vec{E}}_i$

Dimensionality example with GPT-3:

Query and Key projection matrices: $W_Q,W_K\in {\mathbb{R}}^{128\times 12,288}$ — wide matrices, which compress embedding vectors down to key-query-space Token embedding vector $i$: ${\vec{E}}_i\in {\mathbb{R}}^{12,288}$ (column vector) Query and Key vectors $i$: ${\vec{Q}}_i,{\vec{K}}_i\in {\mathbb{R}}^{128}$ — much shorter than the embedding vector (i.e. queries and keys exist in a lower-dimensional space than embeddings)

Query projection — a very wide matrix compresses a tall embedding into a short query vector:

$$\underset{{\vec{Q}}_{i}128\times 1}{\underbrace{\left[\begin{array}{c}{q}_0\\ q_1\\ ⋮\\ q_{127}\end{array}\right]}}=\underset{W_{Q}128\times 12,288}{\underbrace{\left[\begin{array}{cccc}{w}_{0,0}& w_{0,1}& \cdots & w_{0,12287}\\ w_{1,0}& w_{1,1}& \cdots & w_{1,12287}\\ ⋮& ⋮& \ddots & ⋮\\ w_{127,0}& w_{127,1}& \cdots & w_{127,12287}\end{array}\right]}}\underset{{\vec{E}}_{i}12,288\times 1}{\underbrace{\left[\begin{array}{c}{e}_0\\ e_1\\ ⋮\\ ⋮\\ ⋮\\ e_{12287}\end{array}\right]}}$$

Key projection — identical shape, different learned weights:

$$\underset{{\vec{K}}_{i}128\times 1}{\underbrace{\left[\begin{array}{c}{k}_0\\ k_1\\ ⋮\\ k_{127}\end{array}\right]}}=\underset{W_{K}128\times 12,288}{\underbrace{\left[\begin{array}{cccc}{w}_{0,0}^{\prime }& w_{0,1}^{\prime }& \cdots & w_{0,12287}^{\prime }\\ ⋮& ⋮& \ddots & ⋮\\ w_{127,0}^{\prime }& w_{127,1}^{\prime }& \cdots & w_{127,12287}^{\prime }\end{array}\right]}}\underset{{\vec{E}}_{i}12,288\times 1}{\underbrace{\left[\begin{array}{c}{e}_0\\ e_1\\ ⋮\\ ⋮\\ ⋮\\ e_{12287}\end{array}\right]}}$$

The output vector (128-dim) is ~96× shorter than the input (12,288-dim). The projection matrix $W_Q$ has $128\times 12,288\approx 1.57\text{M}$ parameters.

Conceptual framing:

• Query ≈ “what am I looking for?” (e.g. the noun creature looks for preceding adjectives)
• Key ≈ “what do I offer?” (e.g. the adjective fluffy advertises itself as an adjective-in-the-preceding-position)
• Alternatively, $W_Q$ causes each token to ask “questions” (via its query vector) and $W_K$ makes all other tokens attempt to “answer” that question (with their key vectors).

Keep in mind this is an interpretation — what $W_Q$ and $W_K$ actually do is learned, and for most heads is nowhere near this clean.

Note of convenience: $Q_i$ and $K_i$ are defined as column vectors (instead of row vectors)

We have chosen to define the individual token’s query and key vectors, $Q_i$ and $K_i$, as column vectors to look pretty 🙂. Literature commonly defines them as row vectors. A lot of the content of Steps 1-5 may appear transposed in the literature.

As a consequence of defining $Q_i$ and $K_i$ as column vectors,

• Concatenating all $n$ tokens’ query vectors as columns gives the matrix $Q$:

$$Q=\underset{Q128\times n}{\underbrace{\left[\begin{array}{cccc}|& |& & |\\ {\vec{Q}}_0& {\vec{Q}}_1& \cdots & {\vec{Q}}_{n-1}\\ |& |& & |\end{array}\right]}}$$

• Concatenating all $n$ tokens’ key vectors as columns gives the matrix $K$

$$K=\underset{K128\times n}{\underbrace{\left[\begin{array}{cccc}|& |& & |\\ {\vec{K}}_0& {\vec{K}}_1& \cdots & {\vec{K}}_{n-1}\\ |& |& & |\end{array}\right]}}$$

• In the literature (Attention is All You Need), $Q_i$ and $K_i$ are actually defined as row vectors, so the $Q$ and $K$ matrices in the literature are transposed:

$$Q_{\text{literature}}=Q^{\top }=\underset{Q^{\top }n\times 128}{\underbrace{\left[\begin{array}{c}—{\vec{Q}}_0^{\top }—\\ —{\vec{Q}}_1^{\top }—\\ ⋮\\ —{\vec{Q}}_{n-1}^{\top }—\end{array}\right]}},K_{\text{literature}}=K^{\top }=\underset{K^{\top }n\times 128}{\underbrace{\left[\begin{array}{c}—{\vec{K}}_0^{\top }—\\ —{\vec{K}}_1^{\top }—\\ ⋮\\ —{\vec{K}}_{n-1}^{\top }—\end{array}\right]}}$$

Step 2 — Alignment scores

For every (query, key) pair (i.e. all words’ queries against all words’ keys), take the dot product. This tells us how much $E_i$ (the embedding of vector $i$), attends to $E_j$ (the embedding of vector $j$):

${\text{score}}_{ij}={\vec{Q}}_j\cdot {\vec{K}}_i$

Dimensionality example with GPT-3:

• Query vector of token $j$‘s embedding: ${\vec{Q}}_j\in {\mathbb{R}}^{128}$

  • Concatenating all Query vectors (columns) gives $Q$, the matrix of all tokens’ Query vectors

• Key vector of token $i$‘s embedding: ${\vec{K}}_i\in {\mathbb{R}}^{128}$

  • Concatenating all Key vectors (columns) gives $K$, the matrix of all tokens’ Key vectors

• Score: ${\text{score}}_{ij}\in \mathbb{R}$ — a single scalar

• Full score grid: $S\in {\mathbb{R}}^{n\times n}$ where $n$ = context length

• Single dot product (how much token $E_i$ attends to $E_j$)— two 128-dim vectors collapse to one number:

$${\text{score}}_{ij}=\underset{Q_j^{\top }1\times 128}{\underbrace{\left[\begin{array}{cccc}{q}_0& q_1& \cdots & q_{127}\end{array}\right]}}\underset{K_{i}128\times 1}{\underbrace{\left[\begin{array}{c}{k}_0\\ k_1\\ ⋮\\ k_{127}\end{array}\right]}}=\sum _{m=0}^{127}{q}_m{k}_m\in \mathbb{R}$$

• Full score grid — for a $n$-token context, the full score grid is $n\times n$. The attention pattern matrix $A$ will be the same shape:

$$\begin{array}{rl}S=K^{\top }Q& =\underset{K^{\top }n\times 128}{\underbrace{\left[\begin{array}{c}—{\vec{K}}_0^{\top }—\\ —{\vec{K}}_1^{\top }—\\ ⋮\\ —{\vec{K}}_{n-1}^{\top }—\end{array}\right]}}\underset{Q128\times n}{\underbrace{\left[\begin{array}{cccc}|& |& & |\\ {\vec{Q}}_0& {\vec{Q}}_1& \cdots & {\vec{Q}}_{n-1}\\ |& |& & |\end{array}\right]}}\\ & ={\begin{array}{cccccc}& {\vec{E}}_0\stackrel{W_Q}{\to }{\vec{Q}}_0& {\vec{E}}_1\stackrel{W_Q}{\to }{\vec{Q}}_1& {\vec{E}}_2\stackrel{W_Q}{\to }{\vec{Q}}_2& \cdots & {\vec{E}}_{n-1}\stackrel{W_Q}{\to }{\vec{Q}}_{n-1}\\ {\vec{E}}_0\stackrel{W_K}{\to }{\vec{K}}_0& {\vec{K}}_0\cdot {\vec{Q}}_0& {\vec{K}}_0\cdot {\vec{Q}}_1& {\vec{K}}_0\cdot {\vec{Q}}_2& \cdots & {\vec{K}}_0\cdot {\vec{Q}}_{n-1}\\ {\vec{E}}_1\stackrel{W_K}{\to }{\vec{K}}_1& {\vec{K}}_1\cdot {\vec{Q}}_0& {\vec{K}}_1\cdot {\vec{Q}}_1& {\vec{K}}_1\cdot {\vec{Q}}_2& \cdots & {\vec{K}}_1\cdot {\vec{Q}}_{n-1}\\ {\vec{E}}_2\stackrel{W_K}{\to }{\vec{K}}_2& {\vec{K}}_2\cdot {\vec{Q}}_0& {\vec{K}}_2\cdot {\vec{Q}}_1& {\vec{K}}_2\cdot {\vec{Q}}_2& \cdots & {\vec{K}}_2\cdot {\vec{Q}}_{n-1}\\ ⋮& ⋮& ⋮& ⋮& \ddots & ⋮\\ {\vec{E}}_{n-1}\stackrel{W_K}{\to }{\vec{K}}_{n-1}& {\vec{K}}_{n-1}\cdot {\vec{Q}}_0& {\vec{K}}_{n-1}\cdot {\vec{Q}}_1& {\vec{K}}_{n-1}\cdot {\vec{Q}}_2& \cdots & {\vec{K}}_{n-1}\cdot {\vec{Q}}_{n-1}\end{array}}_{n\times n}\end{array}$$

• Notes:
  • Column $j$ holds all scores for token $j$‘s query against every key. This grid scales quadratically with context length — a 2,048-token context produces a $2,048\times 2,048\approx 4.2\text{M}$-entry grid.
  • Column and row labels (${\vec{E}}_0\stackrel{W_Q}{\to }{\vec{Q}}_0$, and ${\vec{E}}_0\stackrel{W_K}{\to }{\vec{K}}_0$) describe Step 1 above

This produces an $n\times n$ grid. Column $j$ contains the scores of every token $i$ as a potential source of information for token $j$. Dot product as an alignment measure: large positive → strongly aligned, zero → unrelated, negative → anti-aligned.

Step 3 — Scale and softmax

Divide by $\sqrt{d_k}$ (this is $\sqrt{128}\approx 11.3$ for GPT-3) for numerical stability (keeps the softmax out of saturated regions) and then apply softmax column by column so each column is a probability distribution:

${\text{attn}}_{ij}={\text{softmax}}_i\left(\frac{{\vec{Q}}_j\cdot {\vec{K}}_i}{\sqrt{d_k}}\right)$

Dimensionality example with GPT-3:

Scale factor: $\sqrt{d_k}=\sqrt{128}\approx 11.3$ Scaled score grid: $S/\sqrt{d_k}\in {\mathbb{R}}^{n\times n}$ — same shape as the raw scores, just divided by 11.3 Attention pattern: $A\in {\mathbb{R}}^{n\times n}$ — same shape, but each column now sums to 1

Take the 5-token phrase: “a big red ball bounced…”

Before softmax — raw scores divided by $\sqrt{128}$ (a 5-token example):

$$\frac{S}{\sqrt{128}}={\left[\begin{array}{ccccc}0.06& -7.40& -2.18& -2.46& -0.46\\ -6.49& 0.26& -0.48& 8.22& -4.26\\ -4.72& -0.50& 0.16& 8.26& -4.91\\ -1.90& -2.63& -4.96& 0.43& -2.86\\ -1.78& -3.62& -7.76& -4.90& 0.05\end{array}\right]}_{5\times 5}$$

After column-wise softmax — each column becomes a probability distribution (sums to 1):

$$A={\left[\begin{array}{ccccc}0.76& 0.00& 0.06& 0.00& 0.36\\ 0.00& 0.65& 0.32& 0.49& 0.01\\ 0.01& 0.30& 0.61& 0.51& 0.00\\ 0.11& 0.04& 0.00& 0.00& 0.03\\ 0.12& 0.01& 0.00& 0.00& 0.60\end{array}\right]}_{5\times 5}\text{where }\sum _i{A}_{ij}=1\forall j$$

The scaling by $\sqrt{128}$ keeps the scores moderate — without it, large dot products push softmax into near-one-hot territory, killing gradients.

The resulting grid is called the attention pattern.

Attention Pattern after the scaling and softmax process (without masking):

5-token phrase: “a big red ball bounced…”

$$\begin{array}{ccccc}\text{Unnormalized}& & \text{Scaled}& & \text{Normalized}\\ \text{Attention Pattern}& & \text{Attention Pattern}& & \text{Attention Pattern}\\ \left[\begin{array}{ccccc}+0.7& -83.7& -24.7& -27.8& -5.2\\ -73.4& +2.9& -5.4& +93.0& -48.2\\ -53.4& -5.7& +1.8& +93.4& -55.6\\ -21.5& -29.7& -56.1& +4.9& -32.4\\ -20.1& -40.9& -87.8& -55.4& +0.6\end{array}\right]& \stackrel{/\sqrt{128}}{\to }& \left[\begin{array}{ccccc}+0.06& -7.40& -2.18& -2.46& -0.46\\ -6.49& +0.26& -0.48& +8.22& -4.26\\ -4.72& -0.50& +0.16& +8.26& -4.91\\ -1.90& -2.63& -4.96& +0.43& -2.86\\ -1.78& -3.62& -7.76& -4.90& +0.05\end{array}\right]& \stackrel{\text{softmax}}{\to }& \left[\begin{array}{ccccc}0.76& 0.00& 0.06& 0.00& 0.36\\ 0.00& 0.65& 0.32& 0.49& 0.01\\ 0.01& 0.30& 0.61& 0.51& 0.00\\ 0.11& 0.04& 0.00& 0.00& 0.03\\ 0.12& 0.01& 0.00& 0.00& 0.60\end{array}\right]\end{array}$$

Step 4 — Masking (autoregressive only)

This step mostly relates to training. Also see the masking section below.

For a next-token-prediction transformer, tokens must not see the future — otherwise training would leak the answer. Before the softmax, set every entry where $i>j$ (future influencing past) to $-\infty$; after softmax those entries become 0. This is called a causal mask.

Masking is what makes it safe to use every position as a training example simultaneously (parallel training). Without it, every position would be “told” what came next.

Attention Pattern after the scaling and softmax process ( WITH masking):

5-token phrase: “a big red ball bounced…”

$$\begin{array}{ccccc}\text{Unnormalized}& & \text{Scaled}& & \text{Normalized}\\ \text{Attention Pattern}& & \text{Attention Pattern}& & \text{Attention Pattern}\\ \left[\begin{array}{ccccc}+0.7& -83.7& -24.7& -27.8& -5.2\\ -\infty & +2.9& -5.4& +93.0& -48.2\\ -\infty & -\infty & +1.8& +93.4& -55.6\\ -\infty & -\infty & -\infty & +4.9& -32.4\\ -\infty & -\infty & -\infty & -\infty & +0.6\end{array}\right]& \stackrel{/\sqrt{128}}{\to }& \left[\begin{array}{ccccc}+0.06& -7.40& -2.18& -2.46& -0.46\\ -\infty & +0.26& -0.48& +8.22& -4.26\\ -\infty & -\infty & +0.16& +8.26& -4.91\\ -\infty & -\infty & -\infty & +0.43& -2.86\\ -\infty & -\infty & -\infty & -\infty & +0.05\end{array}\right]& \stackrel{\text{softmax}}{\to }& \left[\begin{array}{ccccc}1.00& 0.00& 0.06& 0.00& 0.36\\ 0& 1.00& 0.33& 0.49& 0.01\\ 0& 0& 0.62& 0.51& 0.00\\ 0& 0& 0& 0.00& 0.03\\ 0& 0& 0& 0& 0.60\end{array}\right]\end{array}$$

---

We have the attention pattern, describing which words are relevant to (attend to) updating which other words. Now, let’s show how to update all the word embeddings with this new, richer, semantic meaning. It’s what allows “fluffy” in “fluffy creature” to augment the meaning of “creature” by itself.

Step 5 — Values and output

Project every token embedding once more, through the value projection matrix $W_V$, into the embedding space:

${\vec{V}}_i=W_V{\vec{E}}_i$

An interpretation of the value vector, ${\vec{V}}_i$: If the token $i$ is relevant to adjusting the meaning of some other token, what exactly should be added to the embedding of that other token to reflect this?

Dimensionality example with GPT-3 (for one head of attention)

• Value matrix (conceptual): $W_V\in {\mathbb{R}}^{12,288\times 12,288}$ — maps embedding space → embedding space

• Value matrix (in practice, factored): $W_V=V_{\uparrow }{V}_{\downarrow }$ where

  • $V_{\downarrow }\to W_V\in {\mathbb{R}}^{128\times 12,288}$ — maps down: embedding space → 128-dim and
  • $V_{\uparrow }\to W_O\in {\mathbb{R}}^{12,288\times 128}$ — maps up: 128-dim → embedding space
    • For multiple attention heads, all $H$ “up-projections” $V_{\uparrow }^{(h)}\in {\mathbb{R}}^{12,288\times 128}$ are concatenated into one giant matrix, the output matrix $W_O\in {\mathbb{R}}^{12,288\times 128H}$
  • (see low-rank factoring and Terminology gotcha value matrix vs output matrix)

• Token embedding vector $i$: ${\vec{E}}_i\in {\mathbb{R}}^{12,288}$

• Value vector $i$: ${\vec{V}}_i\in {\mathbb{R}}^{12,288}$ — same dimensionality as the embedding

• Conceptual (un-factored) — a square matrix maps a tall vector to an equally tall vector:

$$\underset{V_{i}12,288\times 1}{\underbrace{\left[\begin{array}{c}{v}_0\\ v_1\\ ⋮\\ ⋮\\ ⋮\\ v_{12287}\end{array}\right]}}=\underset{W_{V}12,288\times 12,288}{\underbrace{\left[\begin{array}{ccc}\cdot & \cdots & \cdot \\ ⋮& \ddots & ⋮\\ \cdot & \cdots & \cdot \end{array}\right]}}\underset{E_{i}12,288\times 1}{\underbrace{\left[\begin{array}{c}{e}_0\\ e_1\\ ⋮\\ ⋮\\ ⋮\\ e_{12287}\end{array}\right]}}$$

• Factored (actual) — down-project to 128-dim, then up-project back:

$$\underset{V_{i}12,288\times 1}{\underbrace{\left[\begin{array}{c}{v}_0\\ v_1\\ ⋮\\ ⋮\\ ⋮\\ v_{12287}\end{array}\right]}}=\underset{V_{\uparrow }12,288\times 128}{\underbrace{\left[\begin{array}{ccc}\cdot & \cdots & \cdot \\ ⋮& \ddots & ⋮\\ ⋮& & ⋮\\ \cdot & \cdots & \cdot \end{array}\right]}}\underset{V_{\downarrow }128\times 12,288}{\underbrace{\left[\begin{array}{cccc}\cdot & \cdots & \cdots & \cdot \\ ⋮& \ddots & & ⋮\\ \cdot & \cdots & \cdots & \cdot \end{array}\right]}}\underset{E_{i}12,288\times 1}{\underbrace{\left[\begin{array}{c}{e}_0\\ e_1\\ ⋮\\ ⋮\\ ⋮\\ e_{12287}\end{array}\right]}}$$

• The bottleneck through 128 dims means $W_V$ is low-rank — it can’t represent arbitrary $12,288\times 12,288$ transformations, only ones that pass through a 128-dim intermediate.

The refinement added to token $j$ is the weighted sum of values (${\vec{V}}_i$) weighted by the attention column for $j$ (see Step 3):

$\Delta {\vec{E}}_j={\sum }_i{\text{attn}}_{ij}{\vec{V}}_i$

Dimensionality example with GPT-3:

• Attention weights for token $j$: ${\text{attn}}_{0j},{\text{attn}}_{1j},\cdots \in \mathbb{R}$ — one scalar per token in context, from column $j$ of the attention pattern
• Value vectors: ${\vec{V}}_i\in {\mathbb{R}}^{12,288}$ each
• Refinement vector: $\Delta {\vec{E}}_j\in {\mathbb{R}}^{12,288}$ — same shape as the embedding
• Weighted sum — for token $j$ in an $n$-token-long context, $n$ scalars (from masked (upper triangular) attention pattern) weight $n$ value vectors, producing one vector $\Delta {\vec{E}}_j$:

Example: Take the 5-token context “a big red ball bounced”.

• Let $j=3$, meaning we are interested in updating the token embedding for “ball”
• 5 scalars from Attention (masked) will weight 5 value vectors (one per token), producing one vector $\Delta {\vec{E}}_{j=3}$:
• Recall, the Attention Pattern (“with masking”) above, let’s take its 4th column ($j=3$):

$$A_{\text{masked}}=\left[\begin{array}{ccccc}1.00& 0.00& 0.06& 0.00& 0.36\\ 0& 1.00& 0.33& 0.49& 0.01\\ 0& 0& 0.62& 0.51& 0.00\\ 0& 0& 0& 0.00& 0.03\\ 0& 0& 0& 0& 0.60\end{array}\right]\to A_{\text{masked}}^{ij}=\left[\begin{array}{c}0.00\\ 0.49\\ 0.51\\ 0.00\\ 0\end{array}\right]\forall i,j=3$$

• Then compute $\Delta {\vec{E}}_3$ , the update to “ball“‘s original embedding. Once updated, the new “ball” will include the richer meaning that comes from “big” and “red” which attend to the original “ball”

$$\underset{\Delta {\vec{E}}_{j}12,288\times 1}{\underbrace{\left[\begin{array}{c}\Delta e_0\\ \Delta e_1\\ ⋮\\ \Delta e_{12287}\end{array}\right]}}=\underset{{\text{attn}}_{0j}}{\underbrace{0.00}}\underset{{\vec{V}}_0}{\underbrace{\left[\begin{array}{c}{v}_0^{(0)}\\ v_1^{(0)}\\ ⋮\\ v_{12287}^{(0)}\end{array}\right]}}+\underset{{\text{attn}}_{1j}}{\underbrace{0.49}}\underset{{\vec{V}}_1}{\underbrace{\left[\begin{array}{c}{v}_0^{(1)}\\ v_1^{(1)}\\ ⋮\\ v_{12287}^{(1)}\end{array}\right]}}+\underset{{\text{attn}}_{2j}}{\underbrace{0.51}}\underset{{\vec{V}}_2}{\underbrace{\left[\begin{array}{c}{v}_0^{(2)}\\ v_1^{(2)}\\ ⋮\\ v_{12287}^{(2)}\end{array}\right]}}+\underset{{\text{attn}}_{3j}}{\underbrace{0.00}}\underset{{\vec{V}}_3}{\underbrace{\left[\begin{array}{c}{v}_0^{(3)}\\ v_1^{(3)}\\ ⋮\\ v_{12287}^{(3)}\end{array}\right]}}+\underset{{\text{attn}}_{4j}}{\underbrace{0.00}}\underset{{\vec{V}}_4}{\underbrace{\left[\begin{array}{c}{v}_0^{(4)}\\ v_1^{(4)}\\ ⋮\\ v_{12287}^{(4)}\end{array}\right]}}$$

• Tokens whose attention weight is near 0 contribute almost nothing; tokens near 1 dominate the update. The result is a 12,288-dim vector — a direction in embedding space that gets added to token $j$‘s embedding.

This $\Delta {\vec{E}}_j$ is added (residual-stream style) to the original embedding:

${\vec{E}}_j^{\prime }={\vec{E}}_j+\Delta {\vec{E}}_j$

This attention-weighted update is done for all token embeddings in the context window (not just the $j^{\text{th}}$ token!).

Dimensionality example with GPT-3:

Original embedding: ${\vec{E}}_j\in {\mathbb{R}}^{12,288}$ Refinement from this head: $\Delta {\vec{E}}_j\in {\mathbb{R}}^{12,288}$ Updated embedding: ${\vec{E}}_j^{\prime }\in {\mathbb{R}}^{12,288}$ — same shape, different (updated, richer, more nuanced) direction

Elementwise addition — the original embedding vector gets nudged:

$$\underset{12,288\times 1}{\underbrace{\left[\begin{array}{c}{e}_0^{\prime }\\ e_1^{\prime }\\ ⋮\\ e_{12287}^{\prime }\end{array}\right]}}=\underset{12,288\times 1}{\underbrace{\left[\begin{array}{c}{e}_0\\ e_1\\ ⋮\\ e_{12287}\end{array}\right]}}+\underset{12,288\times 1}{\underbrace{\left[\begin{array}{c}\Delta e_0\\ \Delta e_1\\ ⋮\\ \Delta e_{12287}\end{array}\right]}}$$

This is a residual connection — the identity path means the model can learn to make $\Delta {\vec{E}}_j$ small (or zero) for positions where this head has nothing useful to add.

The compact form

Important

Note: Everything in the literature appears transposed, due to our earlier decision to define a single token’s $Q_i$ and $K_i$ as column vectors (see note above), resulting in the literature’s definition of the $K$ and $Q$ matrices (i.e. for ALL tokens) to differ from ours. Where the literature says $Q{K}^{\top }$, our notes above will say $K^{\top }Q$

The compact notation below, from the paper is an example of this. In our notation, it should instead say $K^{\top }Q$ per our definitions

The whole thing, from the original paper Attention is All You Need:

$\text{Attention}(Q,K,V)=\text{softmax}\left(\frac{Q{K}^{\top }}{\sqrt{d_k}}\right)V$

where $Q,K,V$ are stacks of all query, key, and value vectors across positions. The softmax is understood to apply per column.

Dimensionality example with GPT-3:

Here each token’s vector is a row (the batched convention from the paper). For a context of $n$ tokens: Query matrix: $Q\in {\mathbb{R}}^{n\times 128}$ — each row is one token’s 128-dim query Key matrix: $K\in {\mathbb{R}}^{n\times 128}$ — each row is one token’s 128-dim key Value matrix: $V\in {\mathbb{R}}^{n\times 12,288}$ — each row is one token’s 12,288-dim value

Step 1 — $Q{K}^{\top }$ gives the $n\times n$ score grid:

$$\underset{n\times 128}{\underbrace{\left[\begin{array}{c}—{\vec{q}}_0^{\top }—\\ —{\vec{q}}_1^{\top }—\\ ⋮\\ —{\vec{q}}_{n-1}^{\top }—\end{array}\right]}}\underset{128\times n}{\underbrace{\left[\begin{array}{cccc}|& |& & |\\ {\vec{k}}_0& {\vec{k}}_1& \cdots & {\vec{k}}_{n-1}\\ |& |& & |\end{array}\right]}}=\underset{n\times n}{\underbrace{\left[\begin{array}{cccc}{s}_{00}& s_{01}& \cdots & s_{0,n-1}\\ s_{10}& \ddots & & ⋮\\ ⋮& & \ddots & ⋮\\ s_{n-1,0}& \cdots & \cdots & s_{n-1,n-1}\end{array}\right]}}$$

Step 2 — softmax + multiply by $V$ gives the output:

$$\underset{n\times n}{\underbrace{\text{softmax}\left(\frac{S}{\sqrt{128}}\right)}}\underset{n\times 12,288}{\underbrace{\left[\begin{array}{c}—{\vec{v}}_0^{\top }—\\ —{\vec{v}}_1^{\top }—\\ ⋮\\ —{\vec{v}}_{n-1}^{\top }—\end{array}\right]}}=\underset{n\times 12,288}{\underbrace{\left[\begin{array}{c}—\Delta {\vec{e}}_0^{\top }—\\ —\Delta {\vec{e}}_1^{\top }—\\ ⋮\\ —\Delta {\vec{e}}_{n-1}^{\top }—\end{array}\right]}}$$

The output has one row per token, each row a 12,288-dim refinement vector. The inner dimension (128) cancels in the $Q{K}^{\top }$ product; the inner dimension ($n$) cancels when the attention weights multiply $V$.

Masking and efficient (parallelised) training

Every position in a sequence simultaneously predicts its own next token. So a sequence of length $n$ yields $n$ training signals at no extra cost — provided masking prevents position $j$ from cheating by looking at positions $>j$. This is a huge efficiency multiplier over training on “predict the last token only” one example at a time.

Example: A 12-token-long sequence contains 12 parallel training examples!

• 12-token-long sequence: To| date|,| the| cle|ve|rest| thinker| of| all| time| was…
• Contains the following 12 training examples:
  • To…
  • To| date…
  • To| date|,…
  • To| date|,| the…
  • To| date|,| the| cle…
  • To| date|,| the| cle|ve…
  • To| date|,| the| cle|ve|rest…
  • To| date|,| the| cle|ve|rest| thinker…
  • To| date|,| the| cle|ve|rest| thinker| of…
  • To| date|,| the| cle|ve|rest| thinker| of| all…
  • To| date|,| the| cle|ve|rest| thinker| of| all| time…
  • To| date|,| the| cle|ve|rest| thinker| of| all| time| was…

Also see the masking step above (mostly relevant to training only).

Context size and quadratic cost

The attention pattern has size $n^2$, where $n$ is the context length. Doubling context quadruples attention compute and memory. This is the reason long-context models are nontrivial, and a major research area. update based on state of the art now (Linear Attention, Sliding-Window Attention, State-Space Models, Sparse Attention Mechanisms, Blockwise Attention, Linformer, Reformer, Ring attention, Longformer, Adaptive Attention Span, etc)

Low-rank value factoring

Naïvely, $W_V$ is square ($d_{\text{embed}}\times d_{\text{embed}}$) and would have ~150M parameters in GPT-3 — more than the rest of a single head combined. In practice $W_V$ is factored into two matrices of rank $d_k$: a down-projection into the small key-query space, then an up-projection back to the embedding space. This matches the parameter budget of $W_Q$ and $W_K$ and constrains the value map to be low-rank — see multi-head-attention for why this matters across many heads.

What heads actually do

In the 3b1b running example, one head “has adjectives update nouns”. Real heads do much weirder things. Documented patterns from interpretability work include:

• Previous-token heads (attend strictly to the token immediately before).
• Induction heads (copy patterns — if AB...A appears, predict B).
• Positional heads, duplicate-token heads, anaphora resolution heads, and many more.

Most heads don’t have a clean interpretation at all — they’re just linear algebra that happened to be useful during training.

See also

• self-attention-vs-cross-attention — the two variants of a single attention head
• multi-head-attention — parallel heads in one block
• transformer-architecture — how attention slots into the stack
• softmax — the normalisation used here
• word-embedding, unembedding — the data attention operates on

Sources

• src-3b1b-llms-ch3-attention
• src-3b1b-llms-ch2-transformers