Transformer Basics

2026-09-07 17 min read

Notes from 3Blue1Brown’s chapters on transformers, attention and MLPs, worked through with GPT-3’s numbers throughout. The point is to be able to account for every one of the 175 billion parameters and say what each of them does.

A transformer can be roughly divided into four parts: embedding, attention, MLPs, and unembedding.

The four parts of a transformer: tokens go through the embedding matrix, then 96 layers of attention and MLP writing into a residual stream, then the unembedding matrix and a softmax.

One thing is worth fixing in mind before anything else: attention and the MLP never replace the vector at a position. Each of them computes a change and adds it back. The running vector that everything writes into is usually called the residual stream, and every block reads from it and adds to it.

Embedding

Break up the input into small chunks (i.e. tokens) and turn those chunks into vectors.

The model has a predefined vocabulary of all the tokens it can handle. Tokens are not quite words — they are sub-word pieces produced by byte-level byte-pair encoding, so common words are single tokens and rarer ones get split. The first matrix, called the embedding matrix $\mathbf{W}_E$ , has one column for each token in that vocabulary.

The embedding matrix has one column per token and 12,288 rows; looking up a token means selecting its column.

For GPT-3 the vocabulary has 50,257 tokens and the embedding dimension is 12,288. That gives $12{,}288 \times 50{,}257 = 617{,}558{,}016$ weights.

$\mathbf{W}_E$ starts random and is learned during training. Once it has settled, directions in this high-dimensional space carry semantic meaning — not every direction individually, but many meaningful concepts turn out to correspond to a particular direction, and those directions can be added and subtracted.

Two panels showing that E(king) − E(man) + E(woman) is close to E(queen), and that a single shared direction turns singulars into plurals.

The tool for comparing two vectors is the dot product: large and positive when they point the same way, near zero when they are unrelated, negative when they point opposite ways. Almost everything that follows is built out of dot products.

The number of tokens the model can take in at once is its context size — 2,048 for GPT-3. This is the hard limit on how much text a single forward pass can see.

Unembedding

The desired output is a probability distribution over all the tokens that might come next.

We use the unembedding matrix $\mathbf{W}_U$ , also learned during training. Take the vector sitting at the last position after the final layer, multiply it by $\mathbf{W}_U$ , and the result is a vector with one entry per token in the vocabulary. Then apply softmax to turn it into a probability distribution.

So $\mathbf{W}_U$ has as many rows as the vocabulary size and as many columns as the embedding dimension. For GPT-3 this is another $617{,}558{,}016$ weights.

Only the last position is used at inference time, but during training every position predicts its own next token, so all of them matter.

The softmax function

For a probability distribution we want every element of the vector to lie in $[0, 1]$ and the whole thing to sum to $1$ .

Softmax turns an arbitrary vector into a probability distribution in such a way that its largest value ends up close to $1$ and its smallest values end up close to $0$ .

The process: first raise $e$ to the power of each element, which gives a list of positive values; then normalise those powers of $e$ by their sum.

P_i = \frac{e^{x_i}}{\sum_{n=0}^{N-1}e^{x_n}}

The next token is then sampled from this distribution.

We can introduce a hyperparameter $T$ , called the temperature, to tune the softmax:

P_i = \frac{e^{x_i/T}}{\sum_{n=0}^{N-1}e^{x_n/T}}

Six raw scores turned into probabilities at three temperatures, showing the distribution sharpening as T falls.

When $T$ is larger, lower values get more weight and the distribution becomes more uniform. When $T$ is small, the large values dominate more aggressively. As $T \to 0$ all the weight goes to the maximum. ( $T = 0$ is a limit, not a value you can substitute — the formula divides by $T$ .) In practice most APIs cap $T$ at 2.

Note that because $\exp$ is applied before normalising, differences in the raw scores act multiplicatively: at $T = 1$ , a score two points higher is about $e^2 \approx 7.4$ times more likely.

Attention

The initial embedding is just a table lookup with no context. Each token has one fixed corresponding vector, ignoring position for the moment. It is only in the next step that surrounding tokens get a chance to pass information to each other, so that each token’s vector ends up pointing in a direction that reflects its meaning in this particular sentence.

“Mole” in a mole of carbon, a mole on my skin and whack-a-mole starts as the same vector three times. Attention is what pulls the three apart.

A single head of attention

Query matrix $\mathbf{W}_Q$ . Multiply each token’s embedding by it to get a vector $\vec{Q_i}$ , called a query. This matrix projects the embedding, which is high-dimensional, into a much smaller space usually called the query–key space. The effect can be read as each token asking a question — “is there an adjective in front of me?”

Key matrix $\mathbf{W}_K$ . Multiply each token’s embedding by it to get a vector $\vec{K_j}$ , called a key, projected into the same query–key space. The effect can be read as each token answering — “I am an adjective, and I am in front of you.”

To see which key answers which query, take the dot product $\vec{K_j} \cdot \vec{Q_i}$ . The larger it is, the more the key and query align, and the more relevant token $j$ is to token $i$ . Doing this for every pair gives a grid of scores, one column per query.

Before normalising, the scores are divided by $\sqrt{d_k}$ , where $d_k$ is the dimension of the query–key space ( $\sqrt{128} \approx 11.3$ for GPT-3). The reason is that if the query and key components are independent with zero mean and unit variance, their dot product has standard deviation exactly $\sqrt{d_k}$ . Without the division, that magnitude grows with $d_k$ , the softmax saturates, and the gradients through it become tiny — the original transformer paper’s own justification, offered there as a suspicion rather than a proof. It is a numerical fix, not a conceptual one, but it is always there.

Then each column is normalised with softmax. The result is called the attention pattern.

Masking. In a model that predicts the next token, a token must not be allowed to listen to anything that comes after it — otherwise the answer leaks into the question during training. So before the softmax, every score whose key sits later in the sequence than its query is overwritten with $-\infty$ . Since $e^{-\infty} = 0$ , those entries contribute nothing and the surviving entries in each column still sum to 1.

The 4x4 grid of raw scores with the later-token entries masked to minus infinity, and the same grid after a column-wise softmax.

So far we know how much of each token’s information should flow to each other token. Now we need to actually move it. That takes a third matrix.

Value matrix $\mathbf{W}_V$ . Multiply each token’s embedding by it to get $\vec{V_k}$ , a value. If the naive version is used, this is a square matrix whose dimension equals the embedding dimension. The effect can be read as: if this word turns out to be relevant to something else, what exactly should be added to that something else’s vector to reflect it?

With all the values in hand, the change to each token’s vector is the weighted sum of the values, using the column of the attention pattern belonging to that token.

Here the input and the output of $\mathbf{W}_V$ both live in the embedding space, which would make it an enormous square matrix — $12{,}288 \times 12{,}288 = 150{,}994{,}944$ parameters for a single head. In practice a low-rank factorisation is used instead: the square matrix is written as the product of two rectangular ones, whose shared inner dimension is much smaller than the embedding dimension. Instead of one transformation across the whole embedding space, the vector is first projected down into a small space and then projected back up.

A 12,288 x 12,288 square compared with two thin matrices of shape 128 x 12,288 and 12,288 x 128.

The two halves have standard names: the down-projection is the value matrix proper and the up-projection is the output matrix $\mathbf{W}_O$ . Only the output matrix ever writes into the residual stream.

One head end to end: E splits into queries, keys and values; the queries and keys build the pattern; the values are mixed by it and projected back up.

Counting the parameters. For GPT-3 the query–key space has 128 dimensions. So $\mathbf{W}_Q$ and $\mathbf{W}_K$ each have 12,288 columns and 128 rows, giving $128 \times 12{,}288 = 1{,}572{,}864$ parameters each. Factored the same way, the value and output matrices come to $1{,}572{,}864$ each as well. Adding all four together gives 6,291,456 parameters for one attention head. The naive square value matrix on its own would have been 24 times that.

Multi-headed attention

Multi-headed attention runs several such heads side by side, each with its own query, key, value and output matrices.

For each token, every head proposes a change to be added at that position. All of those proposed changes are summed, one per head, and the total is added to the token’s vector in the residual stream.

The idea is that running many heads in parallel gives the model the capacity to learn many distinct ways that context changes meaning: one head tracking which adjective modifies which noun, another tracking who performed the verb, and so on.

96 head boxes each emitting a proposed change, all summed and added into the residual stream.

GPT-3 has 96 heads in each layer. So one layer’s attention holds $96 \times 6{,}291{,}456 = 603{,}979{,}776$ weights. GPT-3 has 96 such layers, so that number is multiplied by 96 again, reaching 57,982,058,496. That is about one third of GPT-3’s total.

Two footnotes. The 96 output matrices of a layer are usually stitched together into one wide matrix, which is why the whole model comes out to just under 28,000 matrices rather than 37,000. And the same machinery with queries coming from one sequence and keys and values from another is called cross-attention — that is how a translation model reads its source text. Everything above is self-attention, where all three come from the same sequence.

MLP

In one layer of GPT, the multi-headed attention block is followed by an MLP. Two thirds of the parameters actually come from MLPs.

Attention gives the model the ability to fold context into each token’s vector. The MLP is where the facts learned from the training data appear to be stored — where “Michael Jordan” gets connected to “basketball”.

The most important structural difference: the MLP never looks at more than one position at a time. The same two matrices are applied independently to every token’s vector, with no communication between positions at all. Attention is the only place where tokens talk to each other.

The structure

Four steps, applied to each vector on its own:

multiply by an up-projection matrix $\mathbf{W}_\uparrow$ and add a bias $\vec b_\uparrow$ ;
apply a nonlinearity elementwise;
multiply by a down-projection matrix $\mathbf{W}_\downarrow$ and add a bias $\vec b_\downarrow$ ;
add the result back into the residual stream.

The intermediate values, between steps 1 and 3, are the neurons. There are four times as many of them as there are dimensions in the embedding space: $4 \times 12{,}288 = 49{,}152$ for GPT-3.

The nonlinearity used for exposition is the rectified linear unit:

\mathrm{ReLU}(x) = \max(0,\, x)

Negative in, zero out; positive in, unchanged. Real models normally use GELU instead, which has the same shape but is smooth. The GPT line has used it since GPT-1, and GPT-3 inherits it; neither the GPT-2 nor the GPT-3 paper mentions the activation at all, and the released GPT-2 code uses the $\tanh$ approximation of GELU rather than the exact form. Nothing below depends on which one it is.

The MLP: a 12,288-wide vector projected up to 49,152 neurons, passed through ReLU which zeroes about half of them, projected back down, and added to the stream.

What one neuron can do

The useful way to read a matrix multiplication is as a stack of dot products: each row of $\mathbf{W}_\uparrow$ is its own vector, and the corresponding neuron holds the dot product between that row and the vector being processed. So each row is a question the model is asking about this token.

Suppose one row has learned to be the “Michael” direction plus the “Jordan” direction. Then the dot product is roughly 2 for a vector encoding the full name, 1 for a vector encoding only one of the two, and zero or negative otherwise.

Now let the corresponding entry of $\vec b_\uparrow$ be $-1$ . The neuron’s value before the nonlinearity is the dot product minus one — positive if and only if both directions are present. ReLU then clamps everything at or below zero to exactly zero. The neuron fires for Michael Jordan and stays silent for Michael Phelps and Alexis Jordan. The row asks the question, the bias sets the threshold, and ReLU makes it an AND gate.

The other half of the story is $\mathbf{W}_\downarrow$ . Each column of it is a direction in embedding space, and the neuron’s value is the coefficient it is multiplied by. If column $i$ has learned to be a “basketball” direction, then whenever neuron $i$ is on, that direction gets added to the vector; when it is off, nothing happens.

A table showing three input names, their dot product with the row, the bias, the ReLU output, and what gets added.

Rows of $\mathbf{W}_\uparrow$ are the conditions, columns of $\mathbf{W}_\downarrow$ are the payloads, and the neurons are the switches between them. This is very close to the “key–value memory” reading of feed-forward layers in Geva et al. (2021): each neuron’s input weight vector behaves like a key that matches a pattern in the text, and its output weight vector like a value that induces a distribution over what comes next. (Their paper writes both as rows of its own $K$ and $V$ ; whether they land as rows or columns is just a matter of which matrix convention you use.)

Superposition

The clean story above is almost certainly not what really happens. Evidence from interpretability work is that individual neurons rarely stand for one tidy feature; they are polysemantic, firing for several unrelated things.

The reason has to do with geometry. If you insist that features be represented by mutually perpendicular directions, then a $D$ -dimensional space holds exactly $D$ of them. But if you relax “perpendicular” to “nearly perpendicular” — say, within a few degrees — the number you can fit grows exponentially with $D$ . That is a consequence of the Johnson–Lindenstrauss lemma.

And you get most of the way there for free. Take 4,000 random unit vectors and measure the angle of all eight million pairs:

Angle distributions for random unit vectors in 100, 1,000 and 12,288 dimensions; the spread narrows as 1/sqrt(D).

In 100 dimensions a random pair is typically 5.7° off perpendicular and the worst pair is 30° off. In GPT-3’s 12,288 dimensions the typical pair is 0.52° off and every one of the eight million pairs is within 3°. The reason is exact: for two random unit vectors the dot product has mean 0 and variance $1/D$ , and near $90°$ the cosine is locally linear in the angle, so the spread of the angle is $1/\sqrt{D}$ radians. Simply having a big embedding space makes near-orthogonality the default rather than something that has to be arranged.

This means a model can store far more distinct ideas than it has dimensions, which is one plausible reason performance scales so well with size: a space with ten times as many dimensions holds far more than ten times as many independent ideas.

Two caveats worth keeping.

First, “nearly perpendicular” cannot be pushed arbitrarily far. The Welch bound says that for $N$ unit vectors in $D$ dimensions the largest pairwise $|\cos\theta|$ is at least $\sqrt{(N-D)/(D(N-1))}$ , so once $N \gg D$ some pair is always at least about $1/\sqrt{D}$ radians away from perpendicular. Superposition is a statement about the bulk of the pairs, not about all of them.

Second, this is exactly what makes interpretability hard: if features are smeared across neurons rather than sitting one per neuron, reading a single neuron tells you very little. Sparse autoencoders are the best-known tool for trying to pull the underlying features back out, though the picture has moved on: Anthropic’s own circuit-tracing work now uses cross-layer transcoders instead, and how much SAEs buy on downstream tasks is actively contested.

Counting the parameters

$\mathbf{W}_\uparrow$ is $49{,}152 \times 12{,}288 = 603{,}979{,}776$ parameters. $\mathbf{W}_\downarrow$ is the same shape transposed, so the same count. Together that is $1{,}207{,}959{,}552$ per layer, and across 96 layers, 115,964,116,992 — about two thirds of the model. The biases add only 61,440 per layer, which rounds away at this scale.

Putting it together

One layer is: attention over all positions, added into the stream; then an MLP on each position separately, added into the stream. Stack that 96 times, take the vector at the last position, multiply by $\mathbf{W}_U$ , and softmax.

Stacked bar and table of GPT-3's parameter budget: embedding 0.35%, attention 33.1%, MLP 66.2%, unembedding 0.35%.

category	shape	parameters
embedding	$12{,}288 \times 50{,}257$	617,558,016
query	$128 \times 12{,}288 \times 96 \times 96$	14,495,514,624
key	$128 \times 12{,}288 \times 96 \times 96$	14,495,514,624
value	$128 \times 12{,}288 \times 96 \times 96$	14,495,514,624
output	$12{,}288 \times 128 \times 96 \times 96$	14,495,514,624
up-projection	$49{,}152 \times 12{,}288 \times 96$	57,982,058,496
down-projection	$12{,}288 \times 49{,}152 \times 96$	57,982,058,496
unembedding	$50{,}257 \times 12{,}288$	617,558,016
total		175,181,291,520

This is a weights-only count: it leaves out all biases, all layer-norm parameters and the positional embeddings, and it treats the embedding and unembedding as two separate matrices. Counting the real GPT-2-style architecture strictly — biases, two layer norms per block, learned positional embeddings, tied embeddings — gives 174,604,259,328. Both are reconstructions; the paper only ever says “175.0B”.

What this picture leaves out

Worth being explicit about, because none of it appears above:

Positional information. Everything here treats a token’s embedding as independent of where it sits. Real models add positional information — learned embeddings in GPT-3, rotary embeddings in most models since.
Layer normalisation. Every sub-block is wrapped in a normalisation step. It contributes almost no parameters but the model does not train without it.
Attention sparsity. GPT-3 alternates dense layers with locally banded sparse ones; the description above is of a fully dense layer.
Training. Nothing here says where any of these matrices come from. That is the whole of backpropagation and the training objective, which is a separate subject.
Modern variants. Rotary positional embeddings are close to universal; grouped-query attention and mixture-of-experts MLPs are common but not universal; and the MLP is now usually gated — three matrices rather than two, SwiGLU or GeGLU rather than a plain activation. Context windows are orders of magnitude longer. The four-part skeleton survives all of it.

Sources

← More from Research