Transformer Block Assembly in GPT Models

The FeedForward module is a small network of two Linear layers with a GELU activation between them, and it is used in transformer blocks of large language models.

Programming

A Repeated Unit Inside GPT

A GPT model is a large language model built from many repeated transformer blocks. When people say a GPT model has many layers, they are pointing at this repetition. Each transformer block is itself made of smaller parts, and one of those parts is a small neural network called the FeedForward module. This article follows data through that module and then zooms out to show how blocks are assembled into a full model and how the model's output becomes readable text.

The question to keep in mind is a concrete one: if a model is made of one repeated unit, what does that unit actually do to the data passing through it? The FeedForward module is a good place to start, because it carries an important design idea. It temporarily makes the token representations larger so the network can work with a richer space, then brings them back to the size the rest of the model expects.

What do you think happens?

The module is designed for the 124-million-parameter GPT model, where each token is a 768-value embedding. You pass in a batch of 2 samples with 3 tokens each, so the input shape is (2, 3, 768). Before reading any code, what do you think the output shape is? Does the module change the size of the last dimension, and if so, to what?

(2, 3, 768), the same as the input
(2, 3, 3072), because the module expands the embedding
(2, 3, 192), because the module compresses the embedding
(6, 768), because the module flattens the batch and tokens

Reveal answer

Answer:

(2, 3, 768), the same as the input

The input and output dimensions of this module are the same. The embedding dimension is temporarily expanded inside the module and then contracted back. The batch size (2) and the number of tokens (3) are preserved. Only the embedding dimension is changed internally, and it returns to 768 by the end.

Tracing the Embedding Through FeedForward

A common first belief is that the FeedForward module changes the size of its input, so that a 768-value embedding goes in and some different size comes out. The source says otherwise: the input and output dimensions of this module are the same. What does change is what happens between input and output. Inside the module the embedding dimension is expanded and then contracted back.

The module has three steps in order. The first Linear layer maps each token's 768 values up to 3072 values, a factor of 4 larger, since $3072 = 4 \times 768$. A GELU activation then operates on those 3072 values and leaves the shape unchanged. The second Linear layer compresses the 3072 values back into a 768-dimensional representation. The batch size and number of tokens are carried through all three steps untouched.

How the last dimension of the tensor changes at each step: 768 values per token go up to 3072 in the first Linear layer, GELU keeps 3072, and the second Linear layer returns to 768.

Worked example

Tracing one batch through the module

A batch has shape (2, 3, 768): 2 samples, 3 tokens per sample, 768 embedding values per token. Find the shape after each step of FeedForward.

1
Start with the input
The tensor is (2, 3, 768). The last dimension is the embedding size, 768.
2
Apply the first Linear layer
This layer raises the last dimension to 3072. The first two dimensions stay at 2 and 3, so the shape is (2, 3, 3072).
3
Apply GELU
GELU acts on each of the 3072 values. It is a nonlinear activation, but it does not change the shape, so the tensor is still (2, 3, 3072).
4
Apply the second Linear layer
This layer reduces the last dimension from 3072 back to 768. The shape is (2, 3, 768).
5
Compare input and output
The output shape matches the input shape exactly. The module can therefore hand its result to the next part of the model with no reshaping.

Result

The output has shape (2, 3, 768), the same as the input, even though the embedding was temporarily expanded to 3072 values per token.

Reading the FeedForward Class

The FeedForward module is written as a class that inherits from nn.Module. Its constructor builds one nn.Sequential container that holds three steps in order: a Linear layer that maps the embedding size up by a factor of 4, a GELU activation, and a Linear layer that maps the size back down to the original embedding size. The forward method passes the input through that container. The class below is a reconstruction of that design, following the sizes described in the source.

Code walkthrough

import torch
import torch.nn as nn

class GELU(nn.Module):
    def __init__(self):
        super().__init__()
    def forward(self, x):
        return 0.5 * x * (1 + torch.tanh(
            torch.sqrt(torch.tensor(2.0 / torch.pi)) *
            (x + 0.044715 * torch.pow(x, 3))))

class FeedForward(nn.Module):
    def __init__(self, cfg):
        super().__init__()
        self.layers = nn.Sequential(
            nn.Linear(cfg["emb_dim"], 4 * cfg["emb_dim"]),
            GELU(),
            nn.Linear(4 * cfg["emb_dim"], cfg["emb_dim"]),
        )
    def forward(self, x):
        return self.layers(x)

cfg = {"emb_dim": 768}
ffn = FeedForward(cfg)
x = torch.rand(2, 3, 768)
out = ffn(x)
print(out.shape)

1
Line by line: the imports bring in PyTorch and its neural network module.
2
The GELU class defines the activation function applied between the two Linear layers.
3
FeedForward inherits from nn.Module, which is how PyTorch tracks the layers inside it.
4
The constructor reads the embedding size from cfg and builds one nn.Sequential container with three steps in order.
5
The first Linear layer takes emb_dim inputs and produces 4 * emb_dim outputs.
6
The second Linear layer takes 4 * emb_dim inputs and produces emb_dim outputs, closing the expansion.
7
forward simply passes x through the container, so all shape behavior comes from the two Linear sizes.
8
The test creates a random batch of 2 samples with 3 tokens of 768 values each, runs it through the module, and prints the shape.

cfg["emb_dim"] sets the embedding size, 768 for the 124-million-parameter configuration
The first Linear layer maps 768 values up to 4 * 768 = 3072
GELU sits between the two Linear layers and does not change the shape
The second Linear layer maps 3072 values back down to 768
forward does no special work; it passes x through self.layers

Output

torch.Size([2, 3, 768])

Output

torch.Size([2, 3, 768])

The values in out differ on every run because x is random. The shape is fixed by the design: the batch size and token count are preserved, and the last dimension returns to 768.

The parts of the constructor that determine the input and output sizes. The two Linear layers are the only places where shape changes, so a mismatch would be traced back to these arguments.

Best practice

When you read a FeedForward class, look first at the two Linear layers and their size arguments. forward contains no special logic, so the shape of its output is determined entirely by those sizes.

Finding Where a Shape Mismatch Starts

When the output shape disagrees with what you expect, locate the failure by checking the sizes in the order the data passes through. Start with the input shape, then the size given to the first Linear layer, then the size given to the second Linear layer, and finally the output. The first point where the numbers stop matching the design is where the problem is.

Common mistake
Expecting the module to change the embedding size from input to output
What it looks like
Predicting that a (2, 3, 768) input comes out as (2, 3, 3072) because the first layer expands to 3072.
Why it fails
The expansion is internal and temporary. The second Linear layer contracts the values back, so the module's input and output dimensions are the same.
Do this instead
Track the expansion as a middle step. The output shape is (2, 3, 768).
Common mistake
Thinking the batch size or token count changes inside the module
What it looks like
Expecting (2, 3, 768) to become (6, 768) after passing through FeedForward.
Why it fails
The module changes only the embedding dimension. The batch size and number of tokens are preserved.
Do this instead
Keep the first two dimensions fixed and change only the last one internally.
Common mistake
Looking for shape logic inside forward
What it looks like
Searching the forward method for a reshape or a size calculation when the output shape looks wrong.
Why it fails
forward simply passes the input through the nn.Sequential container. All shape behavior comes from the Linear layers and their sizes.
Do this instead
Check the sizes given to the two Linear layers in the constructor.
Common mistake
Mismatching the two Linear sizes so that the second layer does not accept the output of the first, which would break the chain at the point where they connect

How a Transformer Block Is Assembled

Definition

A transformer block is a core structural component of GPT models that combines masked multi-head attention modules with fully connected feed forward networks that use the GELU activation function.

The definition names two parts. The first is the masked multi-head attention module. The second is the fully connected feed forward network, which uses the GELU activation function. Data passes through these parts in order within the block: attention first, then the feed forward network. The FeedForward module you traced above is that second part.

The source places this implementation within a larger plan. The FeedForward module is one of several building blocks needed to assemble the GPT architecture. The section that implements it covers blocks 2 through 4 of that plan: layer normalization, GELU activation, and the feed forward network. Shortcut connections are block 5 and are needed to complete the transformer block.

Stacking Blocks into a GPT Model

A GPT model is a large language model with many repeated transformer blocks, each with millions to billions of parameters in total. Each block has the same two-part structure, and the model is formed by repeating that structure many times. This repetition is what people mean when they say a GPT model has many layers.

The model is a stack of identical transformer blocks, and each block contains an attention part followed by a feed forward part with GELU.

GPT models come in several sizes, measured by parameter count: 124, 345, 762, and 1,542 million parameters. All of them can be implemented with the same GPTModel Python class. The size of a model is therefore set by configuration rather than by writing a different class for each size. The shape-preserving design of FeedForward is part of what makes this possible: because each module returns the dimension it received, blocks can be stacked without reshaping between them.

From Output Tensors to Text

The repeated transformer blocks produce output tensors, and those tensors are numeric. Text generation is the step that decodes them into human-readable text. The source describes this as sequentially predicting one token at a time, based on a given input context. The model does not produce a whole sentence in one step.

On each iteration the model predicts one token from the current context, appends it, and feeds the longer context back in to predict the next token.

The key phrase is one token at a time. Each prediction is based on the input context, and that prediction is what gets decoded. Generation therefore depends on the transformer blocks that produce the output tensors. Because each step extends the context, the next prediction is made from a slightly longer input than the last one.

Why an Untrained Model Produces Noise

Having the right structure is not enough for useful output. Without training, a GPT model generates incoherent text. The architecture defines how data flows through blocks, the FeedForward modules, and the output layer, but the values inside the model's layers determine which tokens it prefers. Before training, those values are not tuned to language, so the decoded tokens do not form coherent text.

Checking Your Understanding of the Shape Path

MEDIUM

A colleague changes the FeedForward constructor so that the second Linear layer maps 3072 values to 1024 instead of 768. Trace a batch of shape (2, 3, 768) through the module. What shape comes out, and at which step does the design stop matching the expected 768-value embedding?

Hints

Walk the shape through the steps in order: input, first Linear, GELU, second Linear, output.
The batch size and token count stay at 2 and 3. Only the last dimension changes.
Compare the final last dimension with the 768 value the rest of the model expects.

Your attempt (not saved)

EASY

Explain in your own words why the first Linear layer's output size and the second Linear layer's input size must agree. Then explain what would happen to the rest of the GPT model if the module's output size did not match its input size.