Building GPT from Scratch: LayerNorm, GELU, Shortcuts, and the First Generated Token

Chapter 3 ended with a working multi-head attention module. It could take a sequence of token representations, decide which earlier positions mattered, and return contextualized vectors. That was a major piece of GPT, but it was still only a piece.

Chapter 4 of Sebastian Raschka's Build a Large Language Model (From Scratch) is where those pieces finally become a model. I went through the official Chapter 4 implementation, rebuilt the path in ordinary Python, and traced the numbers instead of stopping at class names. The result is a complete GPT-style architecture: embeddings, LayerNorm, GELU, feed-forward networks, residual shortcuts, twelve transformer blocks, an output head, and a small greedy generation loop.

The last line is also the important reality check. The architecture can generate text before it has learned anything. With random weights, it produces grammatically shaped nonsense with absolute confidence. Chapter 4 builds the machine; training it is the next problem.

The target: a GPT-2 small-shaped model

The chapter uses a configuration shaped like the smallest GPT-2 model:

GPT_CONFIG_124M = {
    "vocab_size": 50_257,
    "context_length": 1_024,
    "emb_dim": 768,
    "n_heads": 12,
    "n_layers": 12,
    "drop_rate": 0.1,
    "qkv_bias": False,
}

Each value controls a concrete dimension or repeated operation:

Setting	Literal meaning
vocab_size = 50_257	Every token position ends with one score for each GPT-2 vocabulary entry.
context_length = 1_024	The positional embedding table has 1,024 rows, and generation keeps at most that many recent token IDs.
emb_dim = 768	Every token is represented by 768 floating-point features inside the model.
n_heads = 12	The 768 attention features split into 12 heads of 64 features each.
n_layers = 12	The same transformer-block design is instantiated twelve times with different learned weights.
drop_rate = 0.1	During training, dropout randomly removes ten percent of selected activations.
qkv_bias = False	The query, key, and value projections start without bias vectors in this chapter's configuration.

If I pass a batch of two sequences, four tokens each, the shape begins as (2, 4): two rows of token IDs. Token embeddings turn that into (2, 4, 768). Positional embeddings have shape (4, 768) and broadcast across both batch rows when added. Every transformer block preserves (2, 4, 768). The output head changes only the last dimension, producing logits shaped (2, 4, 50257).

That last tensor means: two sequences, four positions per sequence, and 50,257 raw vocabulary scores per position. It is large, but there is no mystery left in the axes.

First build the shell with honest placeholders

Raschka does not begin by implementing every internal operation at once. The chapter first defines a DummyGPTModel whose transformer blocks and final normalization simply return their input. That gives the outer architecture somewhere to stand while the internals are still missing:

class DummyTransformerBlock(nn.Module):
    def forward(self, x):
        return x


class DummyLayerNorm(nn.Module):
    def forward(self, x):
        return x


class DummyGPTModel(nn.Module):
    def __init__(self, cfg):
        super().__init__()
        self.tok_emb = nn.Embedding(cfg["vocab_size"], cfg["emb_dim"])
        self.pos_emb = nn.Embedding(cfg["context_length"], cfg["emb_dim"])
        self.drop_emb = nn.Dropout(cfg["drop_rate"])
        self.trf_blocks = nn.Sequential(
            *[DummyTransformerBlock(cfg) for _ in range(cfg["n_layers"])]
        )
        self.final_norm = DummyLayerNorm(cfg["emb_dim"])
        self.out_head = nn.Linear(
            cfg["emb_dim"], cfg["vocab_size"], bias=False
        )

These placeholders are deliberately boring. They prove that token IDs can enter through the two embedding tables and leave as vocabulary logits with the expected shape. Then each dummy component is replaced by a real one. I like this construction order because it separates two questions that are easy to muddle: “Is the complete pipeline wired correctly?” and “Does this particular layer perform the right mathematics?”

The rest of the chapter answers the second question one component at a time, while the shell keeps the destination visible.

LayerNorm: normalize one token's features

Layer normalization is the first unfamiliar component the chapter builds from scratch. Its job is not to shorten the vector or remove information. Its input and output shapes are identical. It changes the distribution of the values inside each token representation.

I traced the exact example used in my runnable implementation. After a linear layer and ReLU, the first row of activations was:

[0.225952, 0.346954, 0.000000, 0.221604, 0.000000, 0.000000]

The mean is the ordinary average of those six numbers:

mean = 0.132418

The variance measures how far the numbers spread around that mean. Using the population form required here, unbiased=False, I get:

variance = 0.019222

LayerNorm subtracts the mean, then divides by the square root of the variance plus a tiny epsilon. The epsilon is 1e-5; it prevents division by zero when a row has no variation.

normalized = (x - mean) / torch.sqrt(variance + 1e-5)

# first row
[ 0.674460,  1.546987, -0.954852,
  0.643108, -0.954852, -0.954852]

The resulting row has mean 0.0 and variance 0.999480. The variance is very close to one rather than mathematically exact because epsilon is deliberately present.

The complete module is short:

class LayerNorm(nn.Module):
    def __init__(self, emb_dim):
        super().__init__()
        self.eps = 1e-5
        self.scale = nn.Parameter(torch.ones(emb_dim))
        self.shift = nn.Parameter(torch.zeros(emb_dim))

    def forward(self, x):
        mean = x.mean(dim=-1, keepdim=True)
        var = x.var(dim=-1, keepdim=True, unbiased=False)
        norm_x = (x - mean) / torch.sqrt(var + self.eps)
        return self.scale * norm_x + self.shift

The last line matters. Normalization is followed by two learned vectors: scale starts as ones and shift starts as zeros. The network is therefore free to move away from a perfectly standardized distribution when training discovers that a different scale or offset is useful. LayerNorm creates a stable starting coordinate system; it does not permanently handcuff the activations.

The operation runs across dim=-1. For a tensor shaped (batch, tokens, features), every token in every batch item gets its own mean and variance across the feature dimension. This differs from batch normalization, which relies on statistics across examples in a batch. The original Layer Normalization paper was designed precisely to avoid that dependency.

GELU: a smooth activation, not a hard gate

A stack of linear layers without a nonlinear activation would still collapse into one linear transformation. GPT needs a nonlinear function between its feed-forward projections. The chapter implements the Gaussian Error Linear Unit, or GELU, using the common tanh approximation:

class GELU(nn.Module):
    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))
        ))

ReLU is easier to state: negative values become zero and positive values pass unchanged. GELU behaves more gradually. It suppresses negative inputs without chopping every one of them to exactly zero, and it approaches the identity for large positive inputs.

Input	ReLU	GELU
-1.0	0.000000	-0.158808
-0.5	0.000000	-0.154286
0.0	0.000000	0.000000
0.5	0.500000	0.345714
1.0	1.000000	0.841192
3.0	3.000000	2.996363

The curve is not decoration. A smooth derivative gives gradient-based optimization a gradual signal around zero. The GELU paper describes the activation in probabilistic terms, but the implementation lesson is concrete: it multiplies an input by a smooth value between roughly zero and one rather than applying a binary keep-or-delete rule.

The feed-forward network expands before it contracts

Attention mixes information across token positions. The feed-forward network does a different job: it transforms the features of each token independently, using the same weights at every position.

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)

For the 768-dimensional configuration, the first projection expands every token from 768 to 3,072 features. GELU acts on those 3,072 values. The second projection contracts them back to 768.

I also ran a tiny, fully inspectable configuration. Two token vectors with four features each begin as (1, 2, 4). The first linear layer produces (1, 2, 16). The second returns (1, 2, 4). The middle becomes four times wider, while the outer shape is restored so the result can be added to the shortcut path.

This attention-versus-feed-forward distinction is the cleanest way I found to remember the block. Attention lets positions communicate. The feed-forward network gives each position a larger private workspace in which to transform its own feature vector.

Shortcut connections keep a clean path through depth

Deep networks can make gradients shrink as they travel backward through many layers. Residual, or shortcut, connections add the original input of a sublayer back to its output:

shortcut = x
x = layer(x)
x = x + shortcut

The addition requires matching shapes. If x is (2, 4, 768), the transformed branch must also be (2, 4, 768). This is why attention returns to the embedding dimension and why the feed-forward network contracts after expanding.

I tested a five-layer toy network with identical initial weights, once without shortcuts and once with them. After a real backward pass, the mean absolute weight gradients were:

Layer	Plain network	With shortcuts
1	0.00020174	0.22169791
2	0.00012011	0.20694102
3	0.00071520	0.32896996
4	0.00139887	0.26657322
5	0.00504965	1.32585418

This tiny experiment is not a universal benchmark, but it makes the mechanism visible. The shortcut network gives the gradient an additive route that does not require every transformation to preserve it perfectly. The idea was popularized for very deep vision networks by Deep Residual Learning, and it is just as central inside transformers.

One transformer block, in exact execution order

LayerNorm, multi-head attention, GELU, the feed-forward network, dropout, and shortcuts now fit into one block. The order matters:

class TransformerBlock(nn.Module):
    def __init__(self, cfg):
        super().__init__()
        self.att = MultiHeadAttention(
            d_in=cfg["emb_dim"],
            d_out=cfg["emb_dim"],
            context_length=cfg["context_length"],
            num_heads=cfg["n_heads"],
            dropout=cfg["drop_rate"],
            qkv_bias=cfg["qkv_bias"],
        )
        self.ff = FeedForward(cfg)
        self.norm1 = LayerNorm(cfg["emb_dim"])
        self.norm2 = LayerNorm(cfg["emb_dim"])
        self.drop_shortcut = nn.Dropout(cfg["drop_rate"])

    def forward(self, x):
        shortcut = x
        x = self.norm1(x)
        x = self.att(x)
        x = self.drop_shortcut(x)
        x = x + shortcut

        shortcut = x
        x = self.norm2(x)
        x = self.ff(x)
        x = self.drop_shortcut(x)
        x = x + shortcut
        return x

This is a pre-normalization block: LayerNorm runs before attention and before the feed-forward network. Each sublayer has its own normalization and its own residual addition. One transformer block therefore performs two shortcut additions. Twelve blocks perform twenty-four.

The input and output stay shaped (batch, tokens, 768) all the way through. Internally, attention temporarily splits 768 features into twelve 64-feature heads, and the feed-forward network temporarily expands 768 to 3,072. Both return to 768 before their shortcut additions.

Dropout appears on the transformed branches and after the initial embedding sum. It is active during training and disabled by model.eval() during deterministic generation. Dropout is regularization; it is not the causal mask. The causal mask prevents a position from looking into the future. Dropout randomly removes some permitted connections or activations during training.

Stacking the complete GPT model

The full model is surprisingly small as Python source because every large operation is represented by a reusable module:

class GPTModel(nn.Module):
    def __init__(self, cfg):
        super().__init__()
        self.tok_emb = nn.Embedding(cfg["vocab_size"], cfg["emb_dim"])
        self.pos_emb = nn.Embedding(cfg["context_length"], cfg["emb_dim"])
        self.drop_emb = nn.Dropout(cfg["drop_rate"])
        self.trf_blocks = nn.Sequential(
            *[TransformerBlock(cfg) for _ in range(cfg["n_layers"])]
        )
        self.final_norm = LayerNorm(cfg["emb_dim"])
        self.out_head = nn.Linear(
            cfg["emb_dim"], cfg["vocab_size"], bias=False
        )

    def forward(self, in_idx):
        batch_size, seq_len = in_idx.shape
        tok_embeds = self.tok_emb(in_idx)
        pos_embeds = self.pos_emb(
            torch.arange(seq_len, device=in_idx.device)
        )
        x = tok_embeds + pos_embeds
        x = self.drop_emb(x)
        x = self.trf_blocks(x)
        x = self.final_norm(x)
        return self.out_head(x)

The model performs five conceptual steps:

1. Look up one learned 768-value vector for every token ID.
2. Look up one learned 768-value vector for every position and add it to each token vector.
3. Pass the resulting sequence through twelve transformer blocks.
4. Normalize once more after the stack.
5. Project every 768-value token representation to 50,257 vocabulary logits.

The output head does not contain a softmax. It returns raw logits. Training code can feed those logits directly to cross-entropy loss, which applies the necessary normalization stably. Generation code can select or sample from the last position's logits.

Why does the “124M” configuration contain 163 million parameters?

I counted the architecture exactly as written. It has 163,009,536 trainable parameters:

Component	Parameters
Token embedding table	38,597,376
Positional embedding table	786,432
Twelve transformer blocks	85,026,816
Final LayerNorm	1,536
Separate output head	38,597,376
Total as coded	163,009,536

The apparent contradiction lives in the final row before the total. The token embedding table and output head have the same 50257 × 768 shape. GPT-2 ties those weights: the same parameter matrix is used to turn token IDs into embeddings and to turn final hidden vectors back into vocabulary scores. If I count that matrix once rather than twice, the total becomes 124,412,160—the familiar 124M name.

This is a useful lesson in reading model names. “124M” describes the intended GPT-2-small architecture and its tied-weight parameter count. It does not make every educational reimplementation automatically contain exactly 124 million allocated parameters.

The forward pass produces scores, not prose

I seeded PyTorch with 123, created the full untied model, switched it to evaluation mode, and encoded the prompt Hello, I am with the GPT-2 tokenizer. The text becomes four token IDs:

[15496, 11, 314, 716]

After the forward pass, the logits have shape:

(1, 4, 50257)

There is one batch item, four existing token positions, and 50,257 candidate scores at every position. To predict what comes next, the generation loop ignores the first three rows of vocabulary scores and keeps only the final position: logits[:, -1, :]. That leaves (1, 50257).

The highest logit belongs to token ID 27018, which decodes to Feature. The leading space is part of the token. So this randomly initialized network's first greedy continuation is literally:

Hello, I am Feature

That looks promising for exactly one token. It does not survive contact with the next nine:

Hello, I am Featureiman Byeswickattribute argue logger Normandy Compton analogous

This output is not quoted from the book and it is not invented for the article. It is the deterministic output I obtained from the local Chapter 4 implementation with seed 123. The model has architecture but no learned language statistics, so nonsense is the correct result.

Greedy generation is a small loop around a large model

The generation function repeatedly runs the model, selects one token, and appends it:

def generate_text_simple(model, idx, max_new_tokens, context_size):
    for _ in range(max_new_tokens):
        idx_cond = idx[:, -context_size:]

        with torch.no_grad():
            logits = model(idx_cond)

        logits = logits[:, -1, :]
        idx_next = torch.argmax(logits, dim=-1, keepdim=True)
        idx = torch.cat((idx, idx_next), dim=1)

    return idx

Four details make this compact loop worth reading carefully:

1. idx[:, -context_size:] crops the prompt to the model's maximum supported context. Once the sequence exceeds 1,024 tokens, only the most recent 1,024 are passed into this implementation.
2. torch.no_grad() tells PyTorch not to build a backward graph. Generation needs predictions, not gradients.
3. logits[:, -1, :] selects the prediction made at the newest position. Earlier positions predict tokens that are already present.
4. argmax makes decoding greedy and deterministic: always choose the largest score. Applying softmax first would not change which entry is largest, although probabilities become useful when sampling instead.

The selected ID has shape (batch, 1). Concatenating it along dim=1 extends the token sequence by one position. Then the loop starts again with the longer sequence. There is no cache here, so every iteration recomputes every layer for the whole retained context. That is intentionally simple and transparent.

What Chapter 4 actually completes

At the start of this chapter, I had attention. At the end, I have a class whose input is token IDs and whose output is next-token logits. The path is now complete:

token IDs
  → token embeddings + positional embeddings
  → embedding dropout
  → 12 × transformer block
       → LayerNorm → causal multi-head attention → dropout → add shortcut
       → LayerNorm → feed-forward (expand → GELU → contract) → dropout → add shortcut
  → final LayerNorm
  → vocabulary projection
  → logits
  → choose one token and append it

The most important architectural invariant is the stable outer shape. Inside a block, attention splits heads and the feed-forward network expands features. At the boundaries, everything returns to (batch, tokens, 768). That is what allows twelve independently parameterized blocks to compose cleanly and what makes both residual additions legal.

The second important lesson is that a complete architecture is not a trained model. Random weights can pass tensors through every correct shape, produce 50,257 logits, and generate an unlimited stream. None of that means the output is useful. Chapter 4 finishes the mechanism. The next stage is where those 124 million-or-so parameters acquire their values from data.

Still, this is the point where GPT stops looking like one mysterious object. It is token and position lookups, repeated normalization and matrix operations, two additions per block, one final projection, and a loop. Large, yes. Magical, no.