Every large language model is one small block stacked over and over. In this lesson we assemble that block piece by piece and run it ourselves.
Introduction
One Block, Stacked Many Times
Every large language model you have ever used is built from one small machine, stacked over and over, dozens of times. Attention, a little neural network, two skip wires and two normalizations. If you understand this single block, you understand the skeleton of modern language models. Today we build it, count its parameters, and run it.
- Attention + network + wires
- Same block repeated
- Today: build and run it
Where We Left Off
In lesson 55 we derived scaled dot-product attention and multi-head attention with real numbers. Attention lets every token look at other tokens and blend their information. But attention by itself is not a complete layer. It is mostly a weighted averaging machine, and today we add what it is missing and learn how the pieces are wired together.
- Lesson 55: attention with numbers
- Tokens mix information
- Attention alone is incomplete
The Roadmap
Here is the plan. First the feed-forward network and its GELU activation. Then residual connections and layer normalization. We will assemble the block, compare two ways of placing the normalization, and count parameters. Then we meet encoder-only, decoder-only and encoder-decoder models, and finish with a mini GPT block in PyTorch. We do not cover training objectives here.
- Feed-forward, residual, layer norm
- Count parameters
- Three architectures, then code
The Feed-Forward Network
Why Attention Needs Help
Attention moves information between positions, but the mixing itself is just weighted sums of value vectors. There is no rich nonlinear computation on each token. So after attention, every block applies a small two-layer neural network to each token. Think of attention as communication, and this network as private thinking time for each token.
- Attention: communication
- Needs nonlinearity
- Needs per-token computation
The FFN Formula
The formula is: FFN of x equals W two times GELU of W one x plus b one, plus b two. The matrix W one expands the vector from d model to d ff, and W two shrinks it back. A common choice is d ff equal to four times d model. With d model of sixty four, we expand to two hundred fifty six, then compress back to sixty four.
- W1 expands, W2 compresses
- d_ff = 4 x d_model
- Example: 64 to 256 to 64
Why Expand Then Compress?
Why expand? A wider hidden layer gives room for many features to be detected at once. Each hidden unit acts like a pattern detector, and the second matrix writes the results back into the model dimension. This network holds most of the weights. In our example it will turn out to be about two thirds of the parameters in each block.
- More room for features
- Hidden units act like detectors
- About two thirds of block weights
GELU Activation
The activation is GELU. It equals x times capital Phi of x, where Phi is the probability that a standard normal variable is below x. So GELU of one is about zero point eight four, GELU of zero is zero, and GELU of minus one is about minus zero point one six. Unlike ReLU, it is smooth, and small negatives still pass a little gradient.
- GELU(x) = x times Phi(x)
- GELU(1) is about 0.84
- GELU(-1) is about -0.16
Position-Wise: A Common Mistake
A classic mistake is to think the feed-forward network mixes tokens. It does not. It is position-wise: the same two matrices are applied to each token independently. Token three never sees token two inside this layer. All cross-token communication happens in attention, and all per-token computation happens here.
- Same weights at every position
- No mixing between tokens
- Attention talks, FFN computes
Residual Connections
The Skip Connection
Next, the residual connection. Instead of replacing the input with the sublayer's output, we add the two: y equals x plus F of x. Picture a main highway, called the residual stream, running up through the whole model. Each attention or feed-forward sublayer reads from the stream, computes something, and adds its result back.
- y = x + F(x)
- Add the input back
- Residual stream
The Gradient Highway
Why does this help training? Take the derivative. The derivative of y with respect to x is the identity matrix plus the derivative of F. That identity term means there is always a direct path for gradients to flow backward, untouched. Even if the sublayer's own gradient is tiny, the signal still arrives.
- dy/dx = I + dF/dx
- Identity path always exists
- Gradient cannot vanish entirely
Numbers: Twelve Layers
Let us use numbers. Suppose each layer shrinks the gradient by a factor of one half. Across twelve layers without skips, the gradient is zero point five to the twelfth power, about zero point zero zero zero two four four. With skips, the product expands into the identity plus a sum of terms, so a direct path with weight one always survives.
- Scale 0.5 per layer, no skips
- 0.5 to the 12th is 0.000244
- Skips keep an identity path
Shapes Must Match
Residuals demand that the sublayer output has the same shape as its input, which is why d model stays constant through the whole stack. There is also a nice side effect. At initialization, each sublayer adds only a small change, so the block starts close to the identity and learns small corrections rather than rebuilding the representation from scratch.
- Output shape equals input shape
- d_model stays constant
- Blocks start near identity
Layer Normalization
Why Normalize?
Now normalization. As values pass through many layers, their scale can drift: some tokens grow large, others shrink. That makes optimization harder and gradients unstable. Layer normalization rescales each token's feature vector to a consistent mean and spread, so every sublayer sees inputs in a predictable range.
- Activations drift in scale
- Deep stacks amplify drift
- Keep each token stable
The LayerNorm Formula
For one token with d features, compute the mean mu across those features, and the variance sigma squared. Then x hat equals x minus mu, divided by the square root of sigma squared plus epsilon. Epsilon is a tiny constant that avoids dividing by zero. Finally y equals gamma times x hat plus beta, with learnable gamma and beta.
- Mean and variance over features
- x-hat = (x - mu) / sqrt(var + eps)
- y = gamma x-hat + beta
Worked Example
Take the vector two, four, six, eight. The mean is five. The deviations are minus three, minus one, one, three. Squared, they are nine, one, one, nine, whose mean is five. So sigma is the square root of five, about two point two three six. Dividing the deviations gives minus one point three four two, minus zero point four four seven, plus zero point four four seven, plus one point three four two.
- x = [2, 4, 6, 8]
- mu = 5, variance = 5
- x-hat = [-1.34, -0.45, 0.45, 1.34]
Gamma and Beta
Why gamma and beta? Forcing every vector to zero mean and unit variance might remove useful information, so the model can learn its own scale and shift per feature. With gamma equal to two and beta equal to one, our normalized values become minus one point six eight four, zero point one zero six, one point eight nine four, and three point six eight four.
- Learnable scale and shift
- Example: gamma 2, beta 1
- Output: [-1.68, 0.11, 1.89, 3.68]
LayerNorm Versus BatchNorm
In lesson 48, batch norm computed statistics across the batch for each feature. Layer norm flips the axis: it computes statistics across the features of a single token. That means no dependence on other examples, no trouble with padding or variable lengths, small batches work fine, and training and inference behave identically, with no running averages to track.
- BatchNorm: across the batch, per feature
- LayerNorm: across features, per token
- No batch dependence
Mistakes and RMSNorm
A common bug is normalizing over the wrong axis, such as across the sequence. In PyTorch, nn dot LayerNorm of d model normalizes the last dimension, which is the feature axis. You may also hear about RMS norm, a simpler variant used in many modern models that skips subtracting the mean and just rescales by the root mean square.
- Normalize the last axis
- nn.LayerNorm(d_model)
- RMSNorm skips mean subtraction
Assembling the Block
Post-LN: The Original
Now we assemble. The original transformer used post-LN. Run attention, add the input back, then normalize. Run the feed-forward network, add again, normalize again. The normalization sits on the main path, after each residual addition, so the residual stream passes through a normalizer at every sublayer.
- x = LN(x + Attn(x))
- x = LN(x + FFN(x))
- Norm sits after the add
Pre-LN: The Modern Default
Pre-LN moves the normalization inside the branch. Normalize first, then run attention, then add the result to the untouched input. Same for the feed-forward network. Now the residual stream is a clean, unnormalized highway from bottom to top, and the normalizers only act on the side branches that feed the sublayers.
- x = x + Attn(LN(x))
- x = x + FFN(LN(x))
- Clean residual stream
Post-LN Versus Pre-LN
Why do most modern models use pre-LN? Because the clean identity path keeps gradients healthy even in very deep stacks, which makes training more stable and less fragile about the learning rate schedule. Post-LN can work well but is touchier to train at depth. With pre-LN, we add one final layer norm after the last block, before the output layer.
- Pre-LN trains more stably
- Less need for careful warmup
- Add a final LayerNorm
Following the Shapes
Let us follow the shapes through one pre-LN block. The input is batch by time by channels, written B, T, C. Layer norm keeps that shape. Attention returns B, T, C. We add it to the input. Layer norm again, then the feed-forward network expands to B, T, four C inside, and comes back to B, T, C. Add again. The shape never changes.
- Input: batch, time, channels
- Attention keeps shape
- FFN expands inside, same shape out
Counting Parameters
A Small Configuration
Let us count parameters for a small model. Take d model of sixty four, four heads, d ff of two hundred fifty six, two layers, a vocabulary of one thousand tokens, and a context length of one hundred twenty eight. We include biases to be exact. Counting by hand is the best way to build intuition for where the weights live.
- d_model 64, 4 heads
- d_ff 256, 2 layers
- Vocab 1000, context 128
Attention Parameters
Attention has four square matrices: query, key, value, and output projection. Each is sixty four by sixty four, which is four thousand ninety six weights. Four of them make sixteen thousand three hundred eighty four. Add four bias vectors of sixty four each, two hundred fifty six, for a total of sixteen thousand six hundred forty. Splitting into heads does not change the count.
- Four matrices: Q, K, V, output
- Each 64 x 64 = 4,096
- Total with biases: 16,640
Feed-Forward Parameters
The feed-forward network has W one, sixty four by two hundred fifty six, which is sixteen thousand three hundred eighty four weights, plus two hundred fifty six biases. W two is the reverse shape, another sixteen thousand three hundred eighty four, plus sixty four biases. The total is thirty three thousand and eighty eight, roughly twice the attention parameters.
- W1: 64 x 256 + 256 bias
- W2: 256 x 64 + 64 bias
- Total: 33,088
Per-Block Total
Each layer norm has a gamma and a beta of sixty four values, so one hundred twenty eight parameters, and a block has two of them, giving two hundred fifty six. Add everything: sixteen thousand six hundred forty, plus thirty three thousand and eighty eight, plus two hundred fifty six, equals forty nine thousand nine hundred eighty four. A handy rule of thumb is twelve times d model squared, which is forty nine thousand one hundred fifty two.
- Two LayerNorms: 2 x 128 = 256
- 16,640 + 33,088 + 256
- Block = 49,984
The Whole Model
Two blocks give ninety nine thousand nine hundred sixty eight. Token embeddings are one thousand by sixty four, or sixty four thousand. Position embeddings are one hundred twenty eight by sixty four, eight thousand one hundred ninety two. Add the final layer norm, one hundred twenty eight. The total is one hundred seventy two thousand two hundred eighty eight, assuming the output layer reuses the embedding matrix. In big models the blocks dominate, and doubling d model roughly quadruples block parameters.
- Two blocks: 99,968
- Embeddings: 64,000 + 8,192
- Final LN 128, total 172,288
Three Architectures
Same Block, Three Wirings
The block we built is the common ingredient. What changes between model families is the mask and how blocks are wired. There are three main families: encoder-only, like BERT; decoder-only, like GPT; and encoder-decoder, like T five. Each suits different jobs, so let us see how they differ and when to choose each.
- Encoder-only
- Decoder-only
- Encoder-decoder
Encoder-Only: BERT-Style
An encoder-only model uses no causal mask. Every token can attend to every other token, left and right, so each output vector is a rich representation of that token in its full context. This suits understanding tasks: classification, tagging, and building text embeddings for search. It is not a natural text generator, because it was not built to produce tokens one at a time.
- No mask: bidirectional
- Every token sees all tokens
- Classification, tagging, embeddings
Causal Masking
Decoders use a causal mask. Before the softmax, we add minus infinity to every score above the diagonal, so a token cannot see the future. The exponential of minus infinity is zero. Example: token two has scores one and two for tokens one and two. Exponentials are 2.718 and 7.389, summing to 10.107, so the weights are zero point two six nine and zero point seven three one.
- Add minus infinity above diagonal
- Softmax turns it into zero
- Row 2 scores [1, 2] gives 0.27, 0.73
Mask Before Softmax
Here is a common mistake. If you zero out future weights after the softmax, the remaining weights no longer sum to one, and worse, the probabilities were computed while including the future scores. Always mask the scores first, then apply softmax, so the weights over the visible tokens renormalize properly. The mask also lets us train all positions in parallel in one pass.
- Masking after softmax leaks
- Rows no longer sum to one
- Mask first, then normalize
Decoder-Only: GPT-Style
A decoder-only model is a stack of our block with the causal mask. At each position it produces a distribution over the next token, so generation is simply repeating: predict a token, append it, predict again. Because any task can be phrased as text to continue, one architecture handles many tasks through prompting, and this design underlies most general-purpose chat models today.
- Stack of masked blocks
- Predict the next token
- One model, many tasks by prompting
Encoder-Decoder: T5-Style
An encoder-decoder model has two stacks. The encoder reads the whole input bidirectionally. The decoder generates the output with causal masking, but each decoder block has three sublayers: masked self-attention, then cross-attention to the encoder's output, then the feed-forward network. Each one has its own residual connection and normalization.
- Encoder reads the source
- Decoder writes the output
- Decoder block has three sublayers
Cross-Attention and Choosing
In cross-attention, queries come from the decoder, while keys and values come from the encoder output. With three decoder tokens and five source tokens, the attention matrix is three by five, and no causal mask is needed because the whole source is available. This fits translation and summarization. In practice: encoders for understanding, decoders for general generation, encoder-decoders for input to output transformations.
- Queries from decoder, keys and values from encoder
- 3 decoder tokens x 5 source tokens = 3 x 5 matrix
- Pick by task
A Mini GPT Block in PyTorch
The Block Class
Now the code. We define a class Block that inherits from nn dot Module. In init it creates two layer norms, a causal self-attention module, and an MLP made of a linear layer from C to four C, a GELU, and a linear layer back to C. These are exactly the pieces we counted by hand.
- Two LayerNorms
- Multi-head attention
- MLP with GELU
Forward With Shape Comments
The forward method is two lines. First, x equals x plus attention of layer norm one of x. Second, x equals x plus the MLP of layer norm two of x. Inside attention we project to queries, keys and values, reshape to B, heads, T, head dimension, compute scores, apply a lower-triangular mask with minus infinity, softmax, multiply by values, and merge heads back to B, T, C.
- x = x + attn(ln1(x))
- x = x + mlp(ln2(x))
- Every line stays B, T, C
A Tiny Forward Pass
Let us run a tiny example: vocabulary ten, d model eight, two heads of dimension four, batch of two sequences of four tokens. Token ids have shape two by four. Embeddings plus positions give two, four, eight. The block keeps that shape. The final layer norm and output projection give logits of shape two, four, ten: a score for each vocabulary word at each position.
- Tokens: shape (2, 4)
- Embeddings (2, 4, 8)
- Logits (2, 4, 10)
Sanity Checks
Always verify. The attention weights should have shape two, two, four, four. Each row must sum to one, and everything above the diagonal must be exactly zero. Best of all, run a causality test: change the last input token and confirm the outputs at the earlier positions are unchanged. If they move, your mask is leaking the future.
- Attention weights are (2, 2, 4, 4)
- Rows sum to one, upper triangle zero
- Changing a future token changes nothing earlier
Recap and Next Lesson
Recap
Let us recap. A transformer block is attention for communication, a feed-forward network for per-token computation, residual connections as gradient highways, and layer normalization for stable scale. Pre-LN keeps the stream clean. Our small block had forty nine thousand nine hundred eighty four parameters. Encoder-only sees everything, decoder-only uses a causal mask to generate, and encoder-decoder links the two with cross-attention.
- FFN + residual + LayerNorm + attention
- Pre-LN, 49,984 params per block
- Encoder, decoder, encoder-decoder
Next: Pretraining and Fine-Tuning
We built the architecture, but we did not train it. In lesson 57, we see how these models are pretrained, then fine-tuned, including parameter-efficient methods that adapt huge models cheaply.
- Lesson 57 next
- How these models learn
- Subscribe for the series