How Transformer Attention Works, Step by Step

You have probably read that the Transformer is “just attention.” That is the single most misleading thing you can believe about it. If you want to actually understand how transformer attention works — well enough to recognize every piece when you meet it in PyTorch or a diagram — you need to see attention as one station on an assembly line, not the whole factory. This is the builder’s walkthrough: we’ll follow a sentence from raw words to a final prediction, naming each part in the order the data actually flows through it, with just enough code to make it concrete.

This is the technical companion to the plain-English what is attention in AI. If you haven’t met the intuition yet — attention as “which other words should I look at to understand this one?” — start there and come back. Here we open the hood.

How transformer attention works — a vertical assembly line from token embeddings up through positional encoding, attention, and feed-forward layers to a prediction

🔢 Step 1: Words Become Vectors (Embeddings)

A neural network cannot do arithmetic on the word “cat.” So the first thing that happens is tokenization and embedding: the text is split into tokens (roughly words or word-pieces), and each token is looked up in a table that maps it to a list of numbers — a vector. In the original paper this vector has 512 dimensions (called d_model).

That vector is an embedding: a numeric fingerprint that places the token in a space where related meanings sit near each other. Everything downstream — all the attention, all the “understanding” — is arithmetic on these vectors. This is the same idea LIG covers for pictures in image embeddings explained; the Transformer just does it for text and then keeps refining those vectors layer by layer.

At this point our sentence is a stack of vectors, one per token. But there’s an immediate problem to fix before attention can do anything useful.

📍 Step 2: Inject Word Order (Positional Encoding)

Here is a subtle point most explanations skip, and it’s essential. Self-attention, as we’ll see, treats the sentence as an unordered set — every token looks at every other token with no built-in notion of sequence. That means “dog bites man” and “man bites dog” would look identical to it. Disastrous for language.

The fix is positional encoding: before the vectors enter the attention layers, a second pattern of numbers is added to each token’s embedding to stamp it with its position. The original paper uses a fixed pattern of sine and cosine waves of different frequencies — a clever scheme where each position gets a unique, smoothly varying signature. (Modern variants like learned position embeddings, rotary, and ALiBi came later and are not part of the 2017 paper; mention them as “what changed since,” not as the original.)

The takeaway for a builder: order is not free. Attention gives you “which words relate,” but you must hand it “where each word sits” separately. Skip this and the model is word-order-blind.

🎯 Step 3: The Heart — Scaled Dot-Product Attention

Now the main event. For every token, the model builds three different vectors by multiplying its embedding by three learned weight matrices. These are the Query, Key, and Value — Q, K, and V.

The analogy that makes it click: it’s a lookup. Each token issues a Query (“what am I looking for?”), every token advertises a Key (“here’s what I offer”), and carries a Value (“here’s what I’ll hand over if you pick me”). Attention matches queries against keys to decide how much of each value to take.

Mechanically, for a given token:

Score its Query against every token’s Key with a dot product. A big dot product means “these two are relevant to each other.”
Scale every score by dividing by the square root of the key dimension (√d_k). This is the “scaled” in scaled dot-product attention, and it’s not cosmetic: without it, large dot products push the next step into regions where learning stalls. It keeps training stable.
Softmax the scaled scores so they become positive weights that sum to 1 — a distribution over “how much attention to pay to each token.”
Blend: take a weighted sum of every token’s Value using those weights. That weighted sum is the token’s new, context-aware vector.

In PyTorch the whole operation is about four lines:

import torch, torch.nn.functional as F def attention(Q, K, V):    d_k = Q.size(-1)    scores = Q @ K.transpose(-2, -1) / d_k**0.5   # scaled dot products    weights = F.softmax(scores, dim=-1)           # sum to 1    return weights @ V                            # blended values

That’s it. Every mystical phrase you’ve read about attention resolves to: compute similarity between queries and keys, turn it into weights, take a weighted average of values. It is a weighted average — powerful, learnable, but arithmetic, not cognition.

Scaled dot-product attention as a lookup — a Query matched against Keys produces weights that blend the Values into a new context-aware vector

🧠 Step 4: Do It Eight Times (Multi-Head Attention)

One set of Q/K/V learns one kind of relationship. But language has many at once — grammar, reference, tone. So the Transformer runs several attention operations in parallel, each with its own learned Q/K/V matrices. The base model uses 8 heads.

Each head projects the tokens into a smaller subspace, does the scaled dot-product attention above, and produces its own blended output. The eight outputs are concatenated and passed through one more learned matrix to mix them back together. Intuitively, different heads can specialize — one tracking long-distance dependencies, another local grammar.

A myth to resist while you build: the tidy story that each head cleanly learns one human-nameable job. A 2019 study (“Are Sixteen Heads Really Better Than One?”) showed many heads can be pruned away after training with little loss — evidence that multi-head attention is useful redundancy and capacity, not eight indispensable specialists. Build it because it works; don’t over-narrate what each head “means.”

🏗️ Step 5: The Parts Everyone Forgets (FFN, Residuals, LayerNorm)

This is where “the Transformer is just attention” falls apart. After the attention step in each layer, the data flows through several more components that are not optional:

Position-wise feed-forward network (FFN). Every token vector, independently, passes through a small two-layer neural network (in the base model it expands from 512 to 2048 dimensions and back). This is where a surprising amount of the model’s raw capacity lives — the FFN holds roughly two-thirds of the parameters. Attention moves information between tokens; the FFN does heavy processing within each token. Remove it and the model is crippled.
Residual connections. Around both the attention block and the FFN, the input is added to the output (a “skip connection”). This lets information and training signals flow cleanly through many stacked layers without degrading — the same trick that made very deep vision networks like ResNet trainable.
Layer normalization. After each of those additions, a normalization step rescales the numbers to keep them well-behaved, which stabilizes and speeds up training.

So one “Transformer layer” is really: attention → add & normalize → feed-forward → add & normalize. Say it plainly: attention is the headline, but the FFN, residuals, and layer norm are what make the thing actually train and work. “Attention is all you need” means “you can drop the old recurrent network,” not “attention is the only component.”

🧱 Step 6: Stack It, and Mind the Two Halves

A single layer isn’t enough. The base Transformer stacks 6 of these layers in its encoder and 6 in its decoder. Each layer refines the token vectors a little more, building up from surface features to richer, more contextual representations as you go deeper.

And here’s a distinction builders must get right: the original Transformer has two halves.

The encoder reads the whole input at once and builds rich representations of it. Its self-attention is unrestricted — every token sees every other token.
The decoder generates output one token at a time. It uses masked self-attention (each position can only see earlier positions, so it can’t “cheat” by peeking at the future word it’s supposed to predict), plus a second attention step that looks back at the encoder’s output (encoder-decoder attention).

Why this matters: the famous descendants each kept a different half. BERT is encoder-only (great for understanding text). GPT is decoder-only (great for generating text — and the lineage behind ChatGPT). So when someone says “the Transformer,” ask which half they mean; flattening “the Transformer” into “GPT” is a common error. The 2017 paper is the full encoder-decoder, built for translation.

Finally, the decoder’s output vectors pass through one linear layer plus softmax to produce a probability distribution over the vocabulary — the model’s prediction for the next token.

⚠️ The Cost You’re Signing Up For: O(n²)

Before you build one on long inputs, know the bill. Because every token attends to every other token, attention compares all pairs — its cost grows with the square of the sequence length (O(n²)). Double the text and you roughly quadruple the attention compute and memory.

For sentence-length translation (the paper’s use case) that’s cheap. For long documents or big context windows it’s the central bottleneck, and it’s why an entire research industry — sparse attention, linear-attention approximations, and IO-aware methods like FlashAttention — exists to fight it. When you reach for nn.TransformerEncoderLayer and it slows to a crawl on long inputs, this is why. Parallelism and short information paths are the Transformer’s gift; quadratic cost in length is the price.

🛠️ Where You Actually Meet This

You will rarely hand-code the four-line attention function in production. In practice you’ll use torch.nn.TransformerEncoderLayer / TransformerDecoderLayer, or — far more commonly — a pre-built model from the Hugging Face transformers library, where all six steps above are assembled, pretrained, and ready to fine-tune. But every one of those objects is exactly the assembly line you just walked: embed → add position → Q/K/V attention → multi-head → FFN with residual and norm → stack → project to a prediction.

Understanding the line, not just calling it, is what lets you debug a shape mismatch, reason about why long inputs are slow, and know which half of the architecture your model even is. That’s the difference between using a Transformer and understanding one — and it’s the same payoff as building a small vision model by hand in build a CNN image classifier: once you’ve assembled the pieces yourself, the black box stops being black.

❓ Frequently Asked Questions

What are Q, K, and V in attention? Query, Key, and Value — three vectors the model builds for each token from its embedding. The Query asks “what am I looking for,” Keys advertise “what each token offers,” and Values are “what gets handed over.” Attention scores queries against keys, turns the scores into weights, and takes a weighted average of the values.

Why divide by the square root of d_k? To keep the dot-product scores from growing too large as the vectors get bigger. Large scores push the softmax into a region where gradients vanish and training stalls. Scaling by √d_k keeps the numbers in a stable range — hence “scaled” dot-product attention.

Is the Transformer really “just attention”? No. Besides attention, every layer critically uses a position-wise feed-forward network (which holds about two-thirds of the parameters), residual connections, and layer normalization — plus positional encodings at the input. The title means “you no longer need recurrence,” not “attention is the only component.”

What’s the difference between the encoder and decoder? The encoder reads the whole input and builds representations, with every token free to attend to every other. The decoder generates output one token at a time using masked self-attention (it can only look at earlier positions) plus attention over the encoder’s output. BERT kept only the encoder; GPT kept only the decoder.

Why do Transformers get slow on long inputs? Because attention compares every token with every other token, its cost scales with the square of the sequence length (O(n²)). Doubling the input roughly quadruples the work, which is why long context is expensive and why techniques like FlashAttention and sparse attention exist.

📚 Further Reading

What Is “Attention” in AI? — the plain-English intuition behind everything on this page, and the hub cornerstone for the Transformer.
What Is an Embedding? — Q, K, and V are all projections of token embeddings; this explains what those vectors are in the first place.
Build a CNN Image Classifier From Scratch — the vision-side companion to “assemble the pieces yourself,” in the same paper-math-to-runnable-layer style.

📚 Recommended Resources

Disclosure: some links below are Amazon affiliate links. If you buy through them, I earn a small commission — at no extra cost to you — which helps keep this site running.