Transformers from Scratch: How Attention Actually Works

Every large language model you have used is a Transformer. The whole idea fits in one line of maths, and the line most people skip over is the one that matters most: divide by the square root of d_k. I built scaled dot-product attention from scratch in NumPy, ran it on a real sentence, and measured what happens with and without that division. Without it, attention collapses onto a single token. With it, it stays alive. This is Deep Part 1 of the AI Engineer Deep series, and it goes one layer deeper than Part 3 of the original series, which introduced tokens, context, and temperature.

The 30-second answer

Attention answers one question for every token: which other tokens in this sentence should I listen to, and how much? Each token gets three vectors: a query (what am I looking for), a key (what do I offer), and a value (what do I say if picked). Score every query against every key with a dot product, divide by the square root of the key size, run softmax so the scores become weights that sum to 1, and take a weighted sum of the values. That is the whole mechanism. Everything else in the Transformer is scaffolding around it: stacking it 6 times, running 8 copies in parallel, and adding a feed-forward network after each one.

What the paper actually built

"Attention Is All You Need" (Vaswani et al., 2017) replaced recurrence and convolution with attention alone. The base model is small by today's standards, and that is why it is the right model to study: every number is checkable.

PieceBase model

Encoder layers6
Decoder layers6
Model width (d_model)512
Attention heads8
Head size (d_k = d_v)64
Feed-forward width (d_ff)2048
Parameters65 million
EN-DE BLEU (WMT 2014)28.4
EN-FR BLEU (WMT 2014)41.8
Training time3.5 days on 8 GPUs

Two numbers in that table deserve attention of their own. First, 28.4 BLEU beat the best existing results, including ensembles, by over 2 BLEU. Second, the paper notes the model reached state of the art after as little as 12 hours of training on 8 P100 GPUs. The point of the paper was not just quality. It was that attention parallelises: an RNN processes a sentence one token at a time, while attention scores all pairs at once. That is why the same idea scaled from 65 million parameters to GPT-3 at 175 billion, 10 times more than any previous non-sparse language model.

The experiment: attention in 20 lines of NumPy

Theory is cheap, so here is the whole mechanism, runnable. One random vector per token, three learned projections, and the formula from the paper.

import numpy as np rng = np.random.default_rng(7) def softmax(x, axis=-1): x = x - x.max(axis=axis, keepdims=True) e = np.exp(x) return e / e.sum(axis=axis, keepdims=True) def scaled_dot_product_attention(Q, K, V, mask=None): d_k = Q.shape[-1] scores = Q @ K.T / np.sqrt(d_k) if mask is not None: scores = np.where(mask, scores, -np.inf) weights = softmax(scores, axis=-1) return weights, weights @ V tokens = ["the", "animal", "did", "not", "cross", "the", "street", "because", "it", "was", "too", "tired"] X = rng.standard_normal((12, 64)) # one random vector per token Wq = rng.standard_normal((64, 16)) / 8 # 8 = sqrt(64) Wk = rng.standard_normal((64, 16)) / 8 Wv = rng.standard_normal((64, 16)) / 8 Q, K, V = X @ Wq, X @ Wk, X @ Wv weights, out = scaled_dot_product_attention(Q, K, V)

That is it. No library attention, no framework. The softmax check passed on every row: each row of the weight matrix sums to exactly 1.0, which is what makes the output a weighted average of the value vectors instead of a sum that grows with sentence length.

What the weights actually show

Here is the row for the token "it" in "the animal did not cross the street because it was too tired":

TokenWeight from "it"

the0.07
animal0.04
did0.08
not0.04
cross0.04
the0.10
street0.04
because0.13
it0.09
was0.02
too0.05
tired0.29

Read that table honestly and it says something most explainers skip: this model is untrained, and its weights are noise. "It" attends most to "tired" (0.29) and barely notices "animal" (0.04). A trained model learns projection matrices that make "it" attend to "animal", because that is what resolves the pronoun. Attention is a mechanism, not understanding. The understanding lives in the learned Q, K, and V projections, and those only exist after training. If someone shows you attention weights from a random model and claims they prove the model "understands" the sentence, they are showing you noise. The mechanism is exact. The meaning is earned by training.

Why the 1/sqrt(d_k) is not decoration

The paper gives the reason in one paragraph: if the components of q and k are independent random variables with mean 0 and variance 1, their dot product has mean 0 and variance d_k. At the paper's d_k of 64, the raw dot products have a standard deviation of 8. Softmax turns large scores into near-certain weights, and a near-certain softmax has near-zero gradients, so the model stops learning. I measured the effect directly: 2,000 random queries against 32 keys at d_k = 64.

VersionMean max weightMean entropy

Unscaled0.8200.491 nats
Scaled by 1/sqrt(64)0.1653.015 nats

Without the scale, one token grabs 82% of the weight on average and the entropy collapses to 0.49 nats: attention has become a hard pick of a single token. With the scale, the biggest weight is 16.5% on average and the entropy is 3.02 nats, close to the maximum of ln(32) = 3.47 for 32 keys. The division keeps the softmax in the range where it can still spread weight and still learn. It is one character in the formula and it decides whether the model trains at all. That is the kind of detail you only notice when you run the code instead of reading the diagram.

The n-squared bill

Attention scores every token against every other token, so the weight matrix is n by n. I timed one attention call at d_k = 64, median of 5 runs with a warmup, on this machine:

Sequence length (n)Weight matrix sizeMedian time

1280.1 MB0.27 ms
2560.3 MB0.69 ms
5121.0 MB2.42 ms
10244.2 MB9.28 ms
204816.8 MB134.95 ms
409667.1 MB303.96 ms

Two things stand out. First, the matrix itself: at 4,096 tokens, one head needs 67.1 MB just to hold the weights, before the model does anything with them. Multiply by 8 heads and 12 layers and the reason long context is expensive is arithmetic, not mystery. Second, the step from 1,024 to 2,048 jumps far past the 4x that pure n-squared math predicts (about 10x in both runs I did). The matrix stops fitting in cache and memory becomes the bottleneck. Quadratic cost is a memory problem before it is a compute problem, which is exactly the observation FlashAttention is built on. Keep this table in mind when you see a model advertise a 1-million-token context window: someone paid for that matrix, in memory, in time, or in engineering.

The decoder rule: never look ahead

The encoder can look at the whole sentence at once. The decoder cannot, or it would cheat: when it predicts word 9, it must not see word 10. The paper enforces this inside the attention itself by masking out illegal connections, setting their scores to negative infinity before the softmax, so their weights become exactly zero. I ran the same sentence with a causal mask, and the "it" row now only touches its own past:

TokenMasked weight from "it"

the0.12
animal0.06
did0.13
not0.06
cross0.06
the0.15
street0.07
because0.20
it0.15
was, too, tired0.00

Every row still sums to 1.0, and the future is exactly zero. This mask is the entire difference between a model that reads and a model that generates. It is also why generation is slow: the decoder produces one token at a time, and each new token re-attends to everything before it. That cost is what the next deep part on inference, KV cache, and sampling will measure.

Multi-head: eight small attentions, not one big one

One attention head averages everything into a single view, and averaging washes out detail. The paper's fix is to run 8 heads in parallel, each with its own learned projections down to d_k = 64, then concatenate the results and project back to 512. One head might track the subject, another the verb, another the position. I counted the parameters of the paper's base model from its Table 3 sizes (d_model 512, d_ff 2048, 6 layers each side):

PieceParameters

One attention sub-layer (4 x 512^2)1,048,576
Encoder attention, 6 layers6,291,456
Decoder attention, 6 layers x 2 sub-layers12,582,912
All attention18,874,368
Feed-forward networks, 12 layers25,165,824
Attention + feed-forward blocks44,040,192
Paper Table 3 total, base model65,000,000

The rest of the 65 million is embeddings and layer norms. Two things follow. First, attention is 18.9 million of the 65 million: the feed-forward networks are the bigger half of the model, and the "attention is all you need" slogan hides how much plain matrix multiplication sits behind it. Second, the decoder carries twice the attention of the encoder because each decoder layer has two attention sub-layers: self-attention plus attention over the encoder output. When people say the decoder is the expensive half, this table is why.

What came after: RoPE and FlashAttention

The 2017 paper used sinusoidal position encodings, adding a fixed sine wave to each token's vector, and chose them because they might allow the model to extrapolate to sequences longer than the ones seen in training. Two later papers changed how every modern model handles the two hard parts of attention: position and the n-squared matrix.

RoPE, rotary position embedding (Su et al., 2021), rotates the query and key vectors by an angle that depends on their position, so the dot product between two tokens depends on their relative distance. The properties the paper lists are the ones that matter in practice: it works for any sequence length, the influence of a token decays as relative distance grows, and it can be added to linear attention too. It is now the default position encoding in most open models, including the Llama family.

FlashAttention (Dao et al., 2022) attacks the memory wall my timing table showed. It keeps the exact same attention outputs but tiles the computation so the n-by-n matrix never sits in GPU memory: it streams blocks through fast on-chip SRAM instead of reading and writing the whole matrix to slow memory. The paper's measured results: 3x faster GPT-2 training at 1K context, 15% faster end-to-end BERT-large training at 512 tokens against the MLPerf 1.1 speed record, 2.4x faster on the long-range arena at 1K to 4K context, and 0.7 better perplexity on GPT-2. It also made long sequences practical at all: the first Transformers to beat chance on Path-X at 16K tokens (61.4% accuracy) and Path-256 at 64K tokens (63.1%). The lesson is general: when the math is fixed, the win comes from how the data moves, not from a new formula.

My opinionated take

Three things I would tell anyone who wants to work with LLMs instead of just use them. First, implement attention once, from scratch, the way this article did. It takes an afternoon, and after it you will never again be confused by a paper that says "we use 8 heads" or "we mask the future". You will know exactly which matrix changed. Second, never trust an attention heatmap from an untrained or barely trained model. My own run shows "it" attending to "tired" at 0.29 and to "animal" at 0.04. The picture looks convincing and means nothing. Third, respect the n-squared bill when you design anything on top of these models. Part 4 on RAG is mostly about keeping the context you send to the model short for this exact reason, and Part 8 on LLMOps shows what the bill looks like when you ignore it. The models are brilliant. The matrix is not free.

What is next in this series

This is Deep Part 1 of the AI Engineer Deep series, a line-by-line revisit of the original 8-part series with real research and real measurements. Deep Part 2 takes on tokenization and BPE: why "unbelievable" is three tokens, why the same text costs a different number of tokens in different languages, and what I measured when I tokenized real text with a real tokenizer. The original series still stands as the map: Part 3, How LLMs Actually Work is where this topic started, Part 5 on evals is how you check whether any of this is working, and Part 7 on MCP is what you build on top once the model itself makes sense.

Originally published on LogicStack. Run the code examples right in your browser with our free interactive Python and SQL playgrounds — no setup, no sign-up.