Transformers From Scratch

The Transformer (2017) replaced RNNs entirely for sequence tasks. Key insight: you don't need recurrence if you have attention. Process the entire sequence in parallel.

A Transformer block = Multi-Head Attention + Feed Forward + Layer Norm + Residual connections.

Encoder block (BERT-style):
  x → [LayerNorm → MHA → +residual] → [LayerNorm → FFN → +residual]

Decoder block (GPT-style):
  Same but with causal (masked) attention — each token can only see previous tokens.

Positional Encoding: attention has no sense of order. We add position information via sinusoidal encodings (original) or learned embeddings (modern models).

The FFN is 4× wider than d_model: this is where most of the model's "knowledge" is stored. MHA finds relationships, FFN processes them.

1import torch
2import torch.nn as nn
3import math
4
5class FeedForward(nn.Module):
6    def __init__(self, d_model, d_ff=None, dropout=0.1):
7        super().__init__()
8        d_ff = d_ff or d_model * 4
9        self.net = nn.Sequential(
10            nn.Linear(d_model, d_ff),
11            nn.GELU(),
12            nn.Dropout(dropout),
13            nn.Linear(d_ff, d_model),
14            nn.Dropout(dropout),
15        )
16    def forward(self, x): return self.net(x)
17
18class TransformerBlock(nn.Module):
19    def __init__(self, d_model, num_heads, dropout=0.1):
20        super().__init__()
21        self.attn = nn.MultiheadAttention(d_model, num_heads,
22                                          dropout=dropout, batch_first=True)
23        self.ff   = FeedForward(d_model, dropout=dropout)
24        self.ln1  = nn.LayerNorm(d_model)
25        self.ln2  = nn.LayerNorm(d_model)
26        self.drop = nn.Dropout(dropout)
27
28    def forward(self, x, mask=None):
29        # Pre-LayerNorm style (more stable than original post-LN)
30        attn_out, _ = self.attn(self.ln1(x), self.ln1(x), self.ln1(x),
31                                attn_mask=mask)
32        x = x + self.drop(attn_out)    # residual connection
33        x = x + self.drop(self.ff(self.ln2(x)))
34        return x
35
36class GPTStyleModel(nn.Module):
37    """Decoder-only transformer (GPT architecture)"""
38    def __init__(self, vocab_size, d_model, num_heads, num_layers,
39                 max_seq_len, dropout=0.1):
40        super().__init__()
41        self.tok_emb = nn.Embedding(vocab_size, d_model)
42        self.pos_emb = nn.Embedding(max_seq_len, d_model)
43        self.blocks  = nn.ModuleList([
44            TransformerBlock(d_model, num_heads, dropout)
45            for _ in range(num_layers)
46        ])
47        self.ln_f = nn.LayerNorm(d_model)
48        self.head = nn.Linear(d_model, vocab_size, bias=False)
49
50    def forward(self, idx):
51        B, T     = idx.shape
52        positions = torch.arange(T, device=idx.device)
53
54        x = self.tok_emb(idx) + self.pos_emb(positions)  # (B, T, d_model)
55
56        # Causal mask: token i can only attend to positions <= i
57        mask = torch.triu(torch.ones(T, T, device=idx.device), diagonal=1).bool()
58
59        for block in self.blocks:
60            x = block(x, mask=mask)
61
62        x   = self.ln_f(x)
63        return self.head(x)    # (B, T, vocab_size) — logits
64
65# ── Small GPT-style model ─────────────────────────────────────────
66model = GPTStyleModel(
67    vocab_size=50257, d_model=256, num_heads=8,
68    num_layers=4, max_seq_len=512
69)
70tokens = torch.randint(0, 50257, (2, 64))  # (batch=2, seq_len=64)
71logits = model(tokens)
72print(f"Logits: {logits.shape}")            # (2, 64, 50257)
73total  = sum(p.numel() for p in model.parameters())
74print(f"Params: {total/1e6:.1f}M")