How LLMs work · Lecture 6 of 7

MLP layers: where the model stores facts

Attention moves information between tokens. But where is it stored that Michael Jordan plays basketball? We look into the second half of every transformer block. There we meet an old friend from lecture 1: plain dense layers of neurons, which make up two thirds of all of GPT-3.

  • 35 to 45 minutes
  • Prerequisites: lectures 1 to 5
  • 6 interactive demos
01

Where we are: the second half of the block

In lecture 5 you saw how attention moves information between tokens. After attention, the vector of the token “Jordan” already knows that “Michael” comes right before it. But every transformer block has a second part. It is called the MLP (multilayer perceptron). In papers you will also find it as the feed-forward network.

One transformer block. Today we care about the highlighted part. The + sign means the result is added to the vector (the residual connection from lecture 5).

The difference between the two parts is simple. Attention is when the tokens talk to each other. The MLP is when each token thinks alone. Exactly the same small network is applied to every vector, but independently. The vector “Michael” does not see at all what happens to the vector “Jordan”.

Attention: tokens send each other information MLPMLPMLP MLP: same weights, each token separately MichaelJordanplays
Attention connects the token columns. The MLP works inside a single column. All three MLPs are the same network with the same weights.

Why do we need this? Complete the sentence “Michael Jordan plays …”. The word “basketball” is not in the text. The model has to have it stored in its weights. Research suggests that many such facts sit precisely in the MLP layers. In this lecture we build a small MLP that holds exactly such a fact and compute it number by number.

02

An old friend: the dense layers from lecture 1

Remember the neuron from lecture 1. It multiplies its inputs by weights, adds them up, adds a bias and passes the result through an activation function. A whole layer of such neurons can be written as one matrix multiplication: σ(W·a + b). Each row of the matrix W holds the weights of one neuron.

The MLP in a transformer is nothing new. It is exactly two such layers in a row: the same weighted sums, the same biases, the same nonlinearity. The only difference is that they are no longer the whole model. They sit inside every block, alternating with attention.

MLP(x) = W↓ · ReLU( W↑ · x + b↑ ) + b↓ x is the vector of one token. W↑ (up-projection) expands it into a bigger space, ReLU throws away negative values, W↓ (down-projection) brings it back to its original size. GPT uses a very similar function called GELU instead of ReLU.
12,288Input

The vector of one token in GPT-3. That is the size of the embedding and of the residual stream.

49,152Hidden layer

Four times as many neurons. Each one is a perceptron with 12,288 weights and one bias.

12,288Output

Back to the original width, so the result can be added to the token's vector.

From Attention Is All You Need (2017), section 3.3

In addition to attention sub-layers, each of the layers in our encoder and decoder contains a fully connected feed-forward network, which is applied to each position separately and identically. This consists of two linear transformations with a ReLU activation in between.
Formula in the paper: FFN(x) = max(0, xW1 + b1)W2 + b2, with dmodel = 512 and dff = 2,048. Vaswani et al., 2017.

Notice the ratio: 512 → 2,048 is exactly four times. GPT-3 does the same thing at scale, 12,288 → 49,152. The GPT-3 paper says so directly: the feed-forward layer is always four times the size of dmodel. And max(0, …) is nothing but ReLU.

Why xW in the paper and not Wx?

The paper writes the vector as a row and multiplies it by the matrix from the right. In lecture 1 we wrote it as a column and multiplied from the left. It is the same operation, just with the matrix flipped. We stick with W·x, where one row of the matrix belongs to one neuron.

03

Step one: every row asks a question

You know from lecture 4 that directions in embedding space carry meaning. Let's imagine there is a direction meaning “first name Michael”, another for “last name Jordan” and another for “basketball”. To keep it simple, these directions are perpendicular to each other and have length 1.

The vector of the token “Jordan” after attention then looks roughly like this: some of the Michael direction, some of the Jordan direction and a pile of other stuff. How does the MLP ask it whether Michael is in there? With a dot product. If a row of the matrix equals the Michael direction, the dot product tells you how much Michael is in the vector: 1 if it is there, 0 if it is not. Other perpendicular directions contribute nothing at all.

z = r · x + b r is one row of the matrix W↑ (the question), x is the token vector, b is the bias. The row r = Michael + Jordan gives 2 when both names are present, 1 for one and 0 for none. A bias of −1 shifts that to 1, 0 and −1.
Try it

Is the feature in the vector?

The blue arrow is the token vector, made of three directions. The amber arrow is one row of the matrix W↑. The green band is the vector's “shadow” on the row, which is the dot product. On the right you see the whole calculation.

Task: pick the question “Michael + Jordan?” and switch between the three people. Only Michael Jordan gives z > 0. Then move the bias to 0 and see who suddenly “counts” too. Also notice that the “other stuff” slider does not change the result: a perpendicular direction has a zero dot product with the question.

One MLP layer in GPT-3 has 49,152 such questions. Each one is a row of the matrix W↑, and all of them are computed at once by a single matrix-vector multiplication. Just like a layer of the digit network, only much bigger.

04

A whole fact in a small MLP, number by number

Now the whole computation. It is a toy: instead of 12,288 dimensions the vector has just 8 named directions, and instead of 49,152 neurons we have 5. Each number in the vector says how strongly that direction is present. The procedure, however, is exactly the same as in GPT-3:

  1. Input. The vector of one token, as attention left it.
  2. Questions (W↑). Each row takes a dot product with the input and adds a bias. That gives a number z for each neuron.
  3. ReLU. Negative z become zero. Only neurons whose answer is a clear “yes” remain.
  4. Answers (W↓). Each active neuron adds its column of the matrix W↓, scaled by its activation.
  5. Addition. The result is added to the original vector. Nothing gets overwritten.
Try it

A fact inside an MLP, step by step

Choose what the token means and walk through the steps. Hover over or tap any number and you will see below the table how it was computed. Numbers in W↑, b, W↓ and x can be changed with the slider at the bottom.

Task: go through steps 1 to 5 for “Michael Jordan” and then for “Michael Phelps”. Then click “Break N1's bias” and see what suddenly gets added to Michael Phelps. Finally try “Jordan (country)”: which neuron responds, and why that one?

What did you just see? Neuron N1 is an AND gate: it responds only when the vector contains the first name Michael and at the same time the last name Jordan. Its column in the matrix W↓ contains the “basketball” direction and the “Bulls” direction. So the whole fact “Michael Jordan → basketball” is stored in two pieces of numbers: one row of W↑ (when) and one column of W↓ (what to add).

Later layers, and finally the model's last step that picks the next token from the vector (we take it apart in lecture 7), then see a strong “basketball” direction in the vector, and the word “basketball” gets a high probability as the next token.

Key idea

The rows of W↑ are questions, the columns of W↓ are answers. The neuron between them decides whether the answer gets added. It is all plain weighted sums and biases from lecture 1.

05

Why ReLU matters so much

Imagine the MLP without ReLU. A neuron with the row Michael + Jordan and a bias of −1 would return z directly: 1 for Michael Jordan, 0 for Michael Phelps, −1 for text with no names at all. That last one is the problem. Every token without those names would get a negative amount of basketball. A linear function simply cannot say “zero unless both conditions hold”.

ReLU solves this by cutting everything negative down to zero. The values −1, 0, 0, 1 become 0, 0, 0, 1. That is the truth table of AND. It is the same trick as with the logic gates in lecture 1, except that here ReLU provides the threshold instead of a step.

GPT-2 and GPT-3 use GELU instead of ReLU. It is a smoothed version of ReLU: GELU(z) = z·Φ(z), where Φ(z) is the probability that a random number from the standard normal distribution is smaller than z. For large positive z it is almost the same as ReLU, for large negative z it gives almost 0, and around zero it has a slight negative “dip” (its lowest value is about −0.17).

Try it

ReLU as an AND gate, GELU as its smooth sibling

The neuron has the row Michael + Jordan. Switch the names on and off and move the bias (or drag directly in the plot). The table below the plot shows all four combinations at once.

Task: find all bias values for which the neuron computes AND. Why do −1.9 and −1 both work, but −2 does not? Then set the bias to −0.5 and watch AND turn into OR. With GELU, notice that even a “switched off” neuron does not give exactly zero.

Why GELU and not ReLU?

ReLU has a sharp kink at zero and a zero slope for all negative inputs, so no gradient flows through them during training. GELU is smooth and has a small slope even for slightly negative z. In practice large models trained a bit better with it. For our reasoning the difference is small: both functions “let through” positive values and “shut” negative ones.

06

Where most of GPT-3's parameters are

Let's count it the same way as in lecture 1. The up-projection has one row for each of the 49,152 neurons, and each row has 12,288 weights. The down-projection has one column for each neuron, also with 12,288 numbers each.

W↑: 49,152 × 12,288 = 603,979,776    W↓: 12,288 × 49,152 = 603,979,776
one block: 1,207,959,552    × 96 blocks = 115,964,116,992 About 116 billion weights. The biases (49,152 + 12,288 = 61,440 per block, 5,898,240 in total) are negligible next to that. Like 3Blue1Brown, we leave them out of the tally.
Try it

MLP parameter counter

Pick a model or set your own vector width and number of layers. The hidden layer is always four times wider, as in the Transformer and GPT-3 papers.

Task: switch between the three models. Then drag dmodel and watch the number of weights grow with the square: a vector twice as wide means four times as many weights. With the number of layers the growth is only linear.

Two thirds of the model

Of GPT-3's 175 billion parameters, about 116 billion belong to the MLP layers, roughly 66%. Attention has about 58 billion. If the model remembers a huge number of facts about the world somewhere, the MLP layers have the most room for it.

07

The residual stream: one vector everyone writes into

In the Michael Jordan lab the MLP's result was added to its input. That is not a detail. It is how the whole transformer works. Each token has one vector that flows through all 96 blocks. Every attention and every MLP reads it and writes a small change into it:

x ← x + Attention(x)     x ← x + MLP(x) This flow is called the residual stream. Before each read the vector is also normalized (layer norm from lecture 5), but the change is added to the original vector.

So blocks don't replace the vector, they add to it. Whatever an earlier block wrote stays there, and later blocks can use it. A fact can be assembled step by step: one layer figures out who it is about, the next adds the sport, the next the team.

Try it

The residual stream: three blocks, one vector

A toy with five directions and three blocks. Each MLP here has a single AND-type neuron, for example ReLU(Michael + Jordan − 1). Use the switches to turn individual parts off and watch their contribution disappear.

Task: play the vector's flow. Then turn off attention in block 1. The MLP in block 1 no longer knows this is Michael Jordan, and the whole chain breaks: basketball, Bulls and NBA are all missing. Turn it back on and switch off only the MLP in block 2. What remains, and why?

08

Superposition: more features than dimensions

So far we have simplified things. Each feature had its own direction and all directions were exactly perpendicular. But in a 12,288-dimensional space only 12,288 exactly perpendicular directions fit. The things a model knows (names, places, teams, grammar, styles) are far, far more numerous.

The trick is in the word “exactly”. If you allow directions to be only almost perpendicular, say between 89° and 91°, the number of directions that fit grows exponentially with the number of dimensions. Mathematically this is related to the Johnson–Lindenstrauss lemma. In two or three dimensions you can't feel it at all; in high dimensions it is a surprisingly strong effect.

Try it

How many almost perpendicular directions fit?

Your browser generates random unit vectors in N dimensions and computes the angle between every pair. The histogram shows how the angles are distributed.

Task: start at 10 dimensions. The angles are spread all over the plot. Then move N up and watch the histogram squeeze towards 90°. At 100 dimensions set 300 vectors: there are more of them than dimensions, and yet most pairs are close to 90°.

This has a big consequence for the model. Features can be stored as almost perpendicular directions, and the model can pack far more of them into the space than it has dimensions. This is called superposition. The price is a little interference: when “Michael” is in the vector, its dot product with an unrelated question is not exactly 0 but a small number. And here bias and ReLU help again: small random overlaps stay below the threshold and the neuron stays silent.

The same goes for the 49,152 neurons of the hidden layer. A feature is not one neuron but a direction in neuron space. That is why, when researchers look at a single neuron, they often find it responds to several completely unrelated things at once. Such neurons are called polysemantic. It is one of the reasons why “reading a model's mind” is hard. One method for pulling features out of a model is sparse autoencoders. They decompose the activations into a much larger number of directions, of which only a few are active at a time.

An honest note

How exactly large models store facts is still an open research question. The field that studies it is called mechanistic interpretability. The Michael Jordan story is a simplified and plausible picture, not a proven mechanism. Experiments do show, though, that many facts can be located, and even edited, precisely in the MLP layers of the middle blocks.

09

Who wrote those rows and columns?

Nobody. In our labs we set the weights by hand so they would be readable. A real model starts with random numbers. Then it gets text and has to predict the next token. When, for the sentence “Michael Jordan plays …”, it gives the word “basketball” a low probability, backpropagation from lecture 3 computes how each weight should shift a little so it does better next time. It shifts the rows of W↑ and the columns of W↓ alike.

GPT-3 went through roughly 300 billion tokens this way. Facts that repeat often in text get “engraved” into the weights. The model never received an instruction like “remember that Jordan plays basketball”. The arrangement we showed today is just one of the solutions gradient descent can stumble upon.

10

Check what you understood

Six short questions. Nothing is sent anywhere, it is just for you.

1. A row of W↑ is (1, 1, 0), the bias is −1 and the token vector is x = (0.8, 0.9, 0.3). What is the neuron's activation after ReLU?

2. How does an MLP layer differ from attention?

3. A model has dmodel = 1,000 and a hidden layer four times wider. How many weights (without biases) does one of its MLP layers have?

4. What does one column of the matrix W↓ represent?

5. The Michael + Jordan neuron has a bias of −1. What happens if you change the bias to 0?

6. How can more than 12,288 features fit into a 12,288-dimensional space?

11

What you take away

  • An MLP layer is two dense layers from lecture 1: weighted sums, biases and a nonlinearity. In every block it follows attention.
  • Each token separately: the MLP is applied to each vector on its own, with the same weights. The tokens don't talk during it.
  • The rows of W↑ are questions, the columns of W↓ are answers. The bias and ReLU (GELU in GPT) turn a neuron into a threshold, for example an AND gate for “Michael and also Jordan”.
  • The residual stream: the MLP's result is added to the vector. Blocks add to the vector, they don't replace it.
  • 116 billion weights: the MLP layers make up about two thirds of GPT-3's 175 billion parameters.
  • Superposition: there are exponentially many almost perpendicular directions, so a model knows more things than it has dimensions. A single neuron, though, often responds to many unrelated things.

Now you know all the parts: tokens and embeddings, attention, MLP layers and the residual stream. In the last lecture we put them together into the whole of GPT-3, run one sentence through it from the first token to the prediction, and add up all 175,181,291,520 parameters.

Lecture six is behind you

Continue to lecture 7

Sources