Module 2 · Building the Model
Chapter 9 of 17 Run it in your browser

Chapter 9: The Transformer Block¶

Name: Teach Your Computer to Write: Build and Train an LLM from Zero
Author: Truong (Jack) Luu

Where we are in the pipeline Figure 9.1: We combine our pieces into a reusable building block.

After processing information with feed-forward and normalizing the math in Chapter 8, we now have all the ingredients. Multi-head attention lets tokens talk to each other, feed-forward layers let tokens think for themselves, and LayerNorm keeps training stable. In this chapter you will:

Combine these parts into a single Lego brick called a Transformer block.
Add residual connections so deep networks can learn effectively.
Stack multiple blocks to build depth.

Words to Know

Transformer Block: A reusable block of code containing attention, feed-forward, and normalization.
Residual Connection: A shortcut that lets information skip a step, keeping the original signal intact.

Theory¶

Putting the Pieces Together¶

A single Transformer block combines our tools into one unit, and its forward pass is two lines of code:

x = x + self.attn(self.ln1(x))  # "listen, then add to what I know"
x = x + self.ff(self.ln2(x))    # "think, then update what I know"

Notice the x = x + ... pattern, which is what makes deep neural networks work.

Trick 1: The Residual Connection¶

Instead of replacing a token's data with the output of the attention layer, we add the new information to the original data. This is called a residual connection, or skip connection (Figure 9.2).

The data flows down a main highway, taking side trips for attention and feed-forward Figure 9.2: Residual connections provide a direct highway through the block.

To see why this matters, imagine you are in a game of telephone, where each person translates the message and passes it on. After 10 rounds, the original message is usually unrecognizable. Residual connections are like passing the original written message alongside the game of telephone, so even if the spoken message gets mangled, the original is still there.

During training, feedback signals travel backward through the layers to adjust weights, and without residual connections this feedback must pass through every layer in sequence. If each layer distorts the signal slightly, the signal either grows too large or shrinks to nearly nothing. The direct +x highway allows feedback to travel cleanly through dozens of stacked layers.

Trick 2: Pre-Layer Norm¶

Notice that we apply LayerNorm before each major step: attention(LayerNorm(x)).

This is the modern "pre-norm" convention used in GPT models (Figure 9.3), in which we normalize the data first, then do the hard work. It ensures that attention and feed-forward always receive well-behaved numbers, which stabilizes training in the early stages.

Data passes through LayerNorm before entering the Attention layer Figure 9.3: Normalize the data before the hard work.

Stacking Blocks¶

One Transformer block can do a lot, but it is not enough, so we stack multiple blocks in sequence (Figure 9.4). Our model uses 4 blocks.

Each block refines each token's list of numbers a little further, like reading a complex document multiple times. On the first pass you ask "who are the characters?", on the second "what are their motivations?", and on the third "what are the themes?" Earlier blocks capture simple patterns, while later blocks capture deeper meaning, which is why deeper models perform better.

Multiple blocks stacked on top of each other, building deeper understanding Figure 9.4: Stacking blocks allows the model to find complex patterns.

On the big picture map in Figure 9.1, these stacked Transformer blocks form the core engine of our model. The engine is ready to turn the embeddings into lists of numbers that carry each token's context.

Code¶

src/ch08_transformer_block.py (excerpt)
class TransformerBlock(nn.Module):
    def __init__(self):
        super().__init__()
        self.attn = MultiHeadAttention()
        self.ff   = FeedForward()
        # Layer normalization applied before attention and feed-forward
        self.ln1  = nn.LayerNorm(config.n_embd)
        self.ln2  = nn.LayerNorm(config.n_embd)

    def forward(self, x):
        # The + creates a residual connection, adding new info
        x = x + self.attn(self.ln1(x))
        x = x + self.ff(self.ln2(x))
        return x

Run the script to see a block process data and to check the parameter count.

Terminal

$ python src/ch08_transformer_block.py
One TransformerBlock parameters: 197,888
  MultiHeadAttention : 65,664
  FeedForward        : 131,712
  LayerNorms (x2)    : 512

Input  shape: torch.Size([2, 10, 128])
Output shape: torch.Size([2, 10, 128])   (same as input)

--- Residual connection demonstration ---
The input is never lost -- it always flows through.

Original token norm  : 11.470
Attention output norm: 4.599
After residual norm  : 12.903  (combined)

--- Stacking multiple blocks ---
Stacking 4 blocks:
  Params per block: 197,888
  Total params    : 791,552  (4 x 197,888)

Input  shape: torch.Size([2, 10, 128])
Output shape: torch.Size([2, 10, 128])  (unchanged after 4 blocks)

What just happened:

Lines 4 through 8 build a block containing attention, feed-forward, and LayerNorm.
Line 12 shows the residual connection combine the original token data with the attention output. The three norm lines print the overall size of each list of numbers.
We stacked 4 blocks and saw the data flow through smoothly.

Shape Check:

Input to TransformerBlock: [Batch, Time, 128]
Output of TransformerBlock: [Batch, Time, 128]

The block is shape-preserving. The input and output have identical shapes, which is what allows us to stack them endlessly like Lego bricks.

Try It¶

Try It

Open src/ch08_transformer_block.py and change the number of layers in the stack from config.n_layers to 12. Run the script again. Notice how the total parameter count grows, but the output shape remains exactly the same.

In Business

When building software systems (like a house-style writing assistant), you want modular, scalable processes. A Transformer block is the ultimate modular unit. If your assistant is not smart enough, you do not have to invent a new architecture. You stack more blocks and train it longer.

Key Takeaways¶

A Transformer block combines LayerNorm, attention, and feed-forward layers.
Residual connections (x + ...) create a direct highway for data, so deep models can learn effectively.
Pre-norm applies LayerNorm before the hard work, stabilizing the math.
Because blocks preserve the shape of the data, we can stack them modularly.

Check Your Understanding¶

Why do we add the attention output to the original input (x + attention)?
What does "pre-norm" mean?
Why does the Transformer block output exactly the same shape it took in?

Assignment¶

Assignment 9: New numbers, same highway

About 15 minutes. You change one number and write no new code. Work in this chapter's Colab notebook or on your own laptop.

Run it. Run src/ch08_transformer_block.py as it is. Copy four lines: Original token norm, Attention output norm, After residual norm and Total params. Here a norm is one number that says how large a token's list of numbers is overall.
Change one number. Find the line torch.manual_seed(42). The seed fixes the random starting numbers, so every run prints the same values. Change 42 to any whole number you like and run the script again.
Compare. Which of the four lines changed, and which did not? In your new run, which is larger, the original token or the attention output? Is After residual norm closer to the original token or to the attention output?
Explain it. In two or three sentences a manager could follow, say what the three norm lines show about a residual connection: what the block keeps and what it adds.

Hand in: both sets of four lines, your answers to step 3, and your sentences.