Learn LLM

The transformer block

A transformer stacks many identical blocks on top of each other. Each block has two sub-layers: multi-head self-attention and a position-wise feed-forward network. Around each sub-layer is a residual connection followed by layer normalization.

Self-attention: mixes information across the sequence.
Feed-forward network: processes each token independently after mixing.
Residual connections: let gradients skip layers and help deep networks train.
Layer normalization: stabilizes hidden activations.

x' = \text{LayerNorm}(x + \text{SelfAttention}(x))

x'' = \text{LayerNorm}(x' + \text{FFN}(x'))

Pre-norm versus post-norm

The original Transformer applied layer normalization after the residual addition, a layout called post-LN. Modern models such as GPT and LLaMA instead apply layer normalization before the attention or feed-forward sub-layer, called pre-LN. Pre-norm keeps the main path of the network closer to the identity function, which makes it easier to train very deep stacks.

\text{Post-LN}: x' = \text{LayerNorm}(x + \text{SelfAttention}(x))

\text{Pre-LN}: x' = x + \text{SelfAttention}(\text{LayerNorm}(x))

Both forms work, but pre-norm has become the default for large decoder-only language models. It reduces gradient instability early in training, allowing deeper models without careful initialization tricks.

Layer normalization

Layer normalization subtracts the mean and divides by the standard deviation of the activations for each token. This keeps the distribution of inputs to each sub-layer stable, which makes training faster and less sensitive to initialization.

\text{LayerNorm}(x) = \frac{x - \mu}{\sqrt{\sigma^2 + \epsilon}} \cdot \gamma + \beta

Position-wise feed-forward network

After attention mixes tokens, each token is passed independently through the same two-layer network. This is why it is called position-wise: the same MLP operates on every position, but with different inputs because attention has already mixed positions.

\text{FFN}(x) = \max(0, x W_1 + b_1) W_2 + b_2

The inner dimension is usually larger than the model dimension, often by a factor of four. The first projection expands, the ReLU adds non-linearity, and the second projection returns to the model dimension. Some modern variants use SwiGLU or GELU activations instead of ReLU.

Stacking blocks

A typical transformer has 12, 24, or even more identical blocks. The output of one block becomes the input to the next. Each block can refine the representation: lower blocks capture syntax and local patterns, higher blocks capture semantics and long-range structure.

# simplified transformer block
def transformer_block(x):
    # x shape: (seq_len, d_model)
    attn_out = multi_head_self_attention(x)
    x = layer_norm(x + attn_out)
    ffn_out = feed_forward(x)
    x = layer_norm(x + ffn_out)
    return x

Encoder-only, decoder-only, and encoder-decoder

The original transformer used an encoder and decoder for translation. The encoder reads the source sentence with bidirectional attention; the decoder generates the target sentence with causal attention and cross-attention to the encoder.

Cross-attention is identical to self-attention except that the keys and values come from the encoder output while the queries come from the decoder. This lets the decoder attend to the entire source sentence while generating one target token at a time.

\text{CrossAttention}(Q_{dec}, K_{enc}, V_{enc}) = \text{softmax}\left(\frac{Q_{dec} K_{enc}^T}{\sqrt{d_k}}\right) V_{enc}

BERT is encoder-only and learns to fill in masked tokens. GPT is decoder-only and learns to predict the next token. Modern large language models are almost all decoder-only because the same architecture can be used for pre-training and generation.

Why decoder-only?

Decoder-only models unify pre-training and generation. They learn P(next token | previous tokens) with causal attention, which is exactly what you need to generate text left-to-right. No separate architecture is required.

From hidden state to vocabulary

After the final transformer block, each position has a hidden vector. A linear projection maps this vector to a score for every token in the vocabulary. Softmax turns these logits into probabilities.

P(w \mid \text{context}) = \text{softmax}(W_{\text{out}} \cdot h + b_{\text{out}})

This output projection is often tied with the input embedding matrix, meaning W_out is the transpose of the embedding matrix. Tying reduces parameters and improves performance because the same semantic space is used at input and output. The output is a probability distribution over the vocabulary, and the next token is sampled from it.

Transformer block stepper

Walk through one transformer block step by step. Watch how input embeddings flow through self-attention, residual connections, layer normalization, and the feed-forward network.

Transformer block stepper

Step 1 of 6

Input embeddings

Each token is represented as a d_model vector, with positional encoding added.

x = \text{Embedding}(token) + \text{PositionalEncoding}(pos)

Multi-head self-attention

Add & layer norm

Feed-forward network

Add & layer norm

Output projection

Key takeaway

A transformer block alternates mixing through self-attention and processing through feed-forward networks, with residual connections and layer normalization for stability. Decoder-only transformers are the dominant architecture for modern language models.

What comes next

Now that we know how a single block works, the next chapter scales the whole model up. Training an LLM requires data, compute, optimizers, and careful monitoring to turn next-token prediction into a useful base model.

Chapter 10: The Transformer Block