Transformer Block Diagram

A transformer block diagram shows the repeated unit inside every large language model: attention lets each token look at other tokens, a feed-forward network transforms each token on its own, and residual additions plus layer normalisation keep training stable. This sample draws a GPT-style decoder-only block with the attention mechanism opened up on the right. It is a deep-learning architecture diagram, not UML.

1. The Diagram

Follow the blue arrows top to bottom; the yellow lines are the residual stream that bypasses each sub-layer.

1. Input tokens are integer ids; the sequence has length T.
2. Token embedding plus position turns each id into a d_model-wide vector (768 in GPT-2 small) and adds position information.
3. Inside the repeated block, LayerNorm normalises each token's vector first. Normalising before the sub-layer (pre-LN) is what modern LLMs use because deep stacks train more stably.
4. Masked multi-head self-attention mixes information between tokens. The causal mask stops a token from seeing later tokens, which is what makes next-token prediction honest.
5. The first yellow + adds the block's input back in (the residual connection), so each sub-layer only learns a correction.
6. A second LayerNorm feeds the feed-forward network: expand to 4 x d_model, apply GELU, project back.
7. A second residual add closes the block, which repeats N times (12 in GPT-2 small).
8. Final LayerNorm, a linear unembedding to vocabulary size and softmax give next-token probabilities.
9. The right panel shows one head: Q, K and V projections, scores QK^T/√d_k plus mask, softmax, a weighted sum of V, then all heads concatenated and projected by W_o.

2. How to Read It

A dashed frame marked "x N" means everything inside is stacked N times with separate weights. Circles with a plus are additions, not concatenations; concatenation is written out in text where it happens.

3. Variant: Post-LN Encoder Block

The block from the original Attention Is All You Need paper puts normalisation after each residual add ("Add & Norm") and uses unmasked attention because an encoder may look at the whole input. Drawn horizontally, it makes the difference from the pre-LN layout easy to spot: where LayerNorm sits relative to the plus.

4. Pros, Cons, and When to Use

Pros

• Shows the residual stream, the key to why deep transformers train.
• Small enough to memorise and redraw on a whiteboard.
• Opening one sub-layer in a side panel keeps both levels readable.
• Makes pre-LN versus post-LN choices explicit.
• Maps directly onto code modules.

Cons

• Hides tensor shapes per head unless annotated.
• Says nothing about KV caching or inference cost.
• Variants (MoE, RoPE, GQA) need extra notes.
• No standard notation; papers differ in direction and colour.

When to use

• Teaching how LLMs work.
• Documenting a custom model or fine-tuning change.
• Comparing architecture variants.

When not to use

• Explaining the serving stack around an LLM.
• Audiences who only need input and output behaviour.

5. As Diagrams-as-Code

The block as a top-down Mermaid flowchart:

flowchart TB
  T[Input tokens] --> E[Token emb + position]
  subgraph Block[x N layers]
    LN1[LayerNorm] --> MHA[Masked multi-head self-attention] --> A1((+))
    A1 --> LN2[LayerNorm] --> FFN[Feed-forward d to 4d to d, GELU] --> A2((+))
  end
  E --> LN1
  E -. residual .-> A1
  A1 -. residual .-> A2
  A2 --> LNF[Final LayerNorm] --> U[Linear unembed] --> S[Softmax next-token probs]

6. Related

• Architecture Diagrams — box-and-arrow notation reference
• Self-Attention — the attention maths in detail
• Normalization and Dropout — what LayerNorm does
• Autoregressive Models — next-token generation
• RAG Architecture Diagram — sibling sample in this section
• ML Training Pipeline Diagram — sibling sample in this section
• Neural Network Architecture Diagram — sibling sample in this section
• LLM Agent Tool-Call Sequence — sibling sample in this section
• AI & Machine Learning diagrams — section index