The Complete Transformer Architecture

L3-05 covered attention. Now the full picture: how does a Transformer assemble attention into the architecture that’s eaten NLP, vision, and most other modalities?

The Original Diagram (2017)

The “Attention Is All You Need” architecture has two stacks:

Input              Output
  ↓                  ↓
Embedding         Embedding
  ↓                  ↓
+ PosEnc          + PosEnc
  ↓                  ↓
┌────────┐       ┌────────┐
│Encoder │       │Decoder │
│  × N   │       │  × N   │
└───┬────┘       └───┬────┘
    └──────────────→─┤  (cross-attention)
                     ↓
                  Linear + Softmax
                     ↓
                 Output probs

Original use case: machine translation (encoder reads source, decoder writes target). Modern: GPT-class models use decoder only, BERT uses encoder only.

Building Block: One Encoder Layer

Each encoder layer has two sub-blocks:

Input
  ↓
LayerNorm
  ↓
Multi-Head Self-Attention
  ↓
+ residual
  ↓
LayerNorm
  ↓
Feed-Forward (FFN)
  ↓
+ residual
  ↓
Output

Two pieces of magic in every Transformer:

Multi-Head Self-Attention

Each token “looks at” every other token (and itself) and decides what to attend to. See L3-05.

Feed-Forward Network (FFN)

A small 2-layer MLP applied independently to each token:

FFN(x) = max(0, x W_1 + b_1) W_2 + b_2

Typically the hidden dim is 4× the model dim — most of a Transformer’s parameters actually live here, not in attention.

Intuition: attention mixes information between tokens; FFN does per-token processing.

Why LayerNorm and Residuals

Residuals (+ x)

Make 12 / 24 / 96 layer networks trainable. Without them, gradient vanishes. See L3-03.

LayerNorm

Normalizes the activation distribution at each layer to keep training stable.

Detail: original paper put LayerNorm after sub-blocks (“Post-LN”). Modern: almost everyone uses Pre-LN (before sub-blocks) — more stable training.

# Pre-LN (modern)
y = x + Attention(LayerNorm(x))
y = y + FFN(LayerNorm(y))

Building Block: One Decoder Layer

Same structure, but with two attention layers:

LayerNorm
  ↓
Masked Multi-Head Self-Attention  ← look at own previous tokens only
  ↓
+ residual
  ↓
LayerNorm
  ↓
Multi-Head Cross-Attention  ← look at encoder output
  ↓
+ residual
  ↓
LayerNorm
  ↓
Feed-Forward (FFN)
  ↓
+ residual

Two attention layers:

1. Self-attention (masked) — see only earlier tokens (causal)
2. Cross-attention — Q from decoder, K/V from encoder

Three Modern Architectural Branches

The original encoder-decoder architecture has evolved into three branches:

1. Encoder Only (BERT family)

• Drop the decoder, keep only the encoder stack
• Train with Masked Language Modeling (MLM)
• Output: a vector for each token, used for understanding tasks
• Examples: BERT, RoBERTa, DeBERTa, ModernBERT

2. Decoder Only (GPT family)

• Drop the encoder, keep only the decoder (without cross-attention)
• Train with next-token prediction
• Output: probability distribution for the next token
• Examples: GPT-2/3/4, LLaMA, Claude, Mistral, DeepSeek, most modern LLMs

3. Encoder-Decoder (T5, BART)

• Keep both
• Train with denoising autoencoding
• Examples: T5, BART, mT5, FLAN-T5

By 2024, decoder-only dominates — its simplicity + scalability is just nicer.

A Real GPT-2 Layer Diagram

        ┌─────────────────────────┐
        │      Token Embedding     │
        └────────────┬─────────────┘
                     │
        ┌────────────▼─────────────┐
        │   + Positional Encoding  │
        └────────────┬─────────────┘
                     │
        ┌────────────▼─────────────┐
        │   Transformer Block 1    │
        │  ┌────────────────────┐  │
        │  │      LayerNorm     │  │
        │  │ Masked Self-Attn   │  │
        │  │      + residual    │  │
        │  │      LayerNorm     │  │
        │  │       FFN (4×d)    │  │
        │  │      + residual    │  │
        │  └────────────────────┘  │
        └────────────┬─────────────┘
                     ↓
                  ... × 12, 24, 96 ...
                     ↓
        ┌────────────▼─────────────┐
        │      Final LayerNorm     │
        └────────────┬─────────────┘
                     │
        ┌────────────▼─────────────┐
        │     Linear + Softmax     │
        └────────────┬─────────────┘
                     │
                 Next token
                  probabilities

Hyperparameter Conventions

Model	Layers	d_model	Heads	FFN
GPT-2 small	12	768	12	3072 (4×)
GPT-3	96	12288	96	49152 (4×)
BERT-base	12	768	12	3072
LLaMA-7B	32	4096	32	11008 (~2.7×)
LLaMA-70B	80	8192	64	28672

You’ll notice: as models get bigger, both depth (more layers) and width (d_model) grow.

Build a Mini Transformer

A complete GPT-2 in ~50 lines:

import torch
import torch.nn as nn

class TransformerBlock(nn.Module):
    def __init__(self, d_model, n_heads, ffn_dim):
        super().__init__()
        self.ln1 = nn.LayerNorm(d_model)
        self.attn = nn.MultiheadAttention(d_model, n_heads, batch_first=True)
        self.ln2 = nn.LayerNorm(d_model)
        self.ffn = nn.Sequential(
            nn.Linear(d_model, ffn_dim),
            nn.GELU(),
            nn.Linear(ffn_dim, d_model),
        )

    def forward(self, x, mask=None):
        # Pre-LN
        h = self.ln1(x)
        attn_out, _ = self.attn(h, h, h, attn_mask=mask)
        x = x + attn_out
        x = x + self.ffn(self.ln2(x))
        return x

class GPT(nn.Module):
    def __init__(self, vocab_size, d_model=768, n_layers=12, n_heads=12, max_len=1024):
        super().__init__()
        self.tok_emb = nn.Embedding(vocab_size, d_model)
        self.pos_emb = nn.Embedding(max_len, d_model)
        self.blocks = nn.ModuleList([
            TransformerBlock(d_model, n_heads, 4 * d_model)
            for _ in range(n_layers)
        ])
        self.ln_f = nn.LayerNorm(d_model)
        self.head = nn.Linear(d_model, vocab_size, bias=False)

    def forward(self, idx):
        B, L = idx.shape
        pos = torch.arange(L, device=idx.device).unsqueeze(0)
        x = self.tok_emb(idx) + self.pos_emb(pos)

        # Causal mask
        mask = torch.triu(torch.ones(L, L, device=idx.device), diagonal=1).bool()

        for block in self.blocks:
            x = block(x, mask)

        x = self.ln_f(x)
        return self.head(x)  # logits over vocab

That’s it. All the magic of modern LLMs lives in this ~50-line skeleton plus enormous scale and good training.

Why Transformer Beat RNN/CNN

Dimension	RNN/LSTM	CNN	Transformer
Parallelization	Sequential, slow	Parallel	Parallel
Long deps	Hard	Hard	Easy
Inductive bias	Sequence	Local + translation	Almost none
Param efficiency	High (per param)	High	Medium
Scales to massive data	Worse	Worse	Excellent

The decisive thing: Transformers benefit more from scale than alternatives. Combined with abundant GPU + abundant data — they win.

Next recommended: L3-09 BERT vs GPT or L4-01 LLM Training.