Transformers From Scratch: What I Wish I Knew on Day 1
A minimal, math-light walkthrough of attention, multi-head layers, and positional encoding — with clean PyTorch code you can actually read.
Why another Transformer post?
Most tutorials either drown you in equations or hand you a black-box nn.Transformer. This is the post I wish someone had handed me on day 1 — minimal math, real code, and an honest map of what actually matters.
The one-sentence intuition
A Transformer is a learnable lookup: every token asks every other token "how relevant are you to me?" and averages them accordingly.
That's it. Everything else — multi-head, positional encoding, feed-forward blocks — is engineering polish on top of that idea.
Scaled dot-product attention
Given queries Q, keys K, values V, attention is:
import torch
import torch.nn.functional as F
def attention(Q, K, V, mask=None):
d_k = Q.size(-1)
scores = Q @ K.transpose(-2, -1) / (d_k ** 0.5)
if mask is not None:
scores = scores.masked_fill(mask == 0, float('-inf'))
attn = F.softmax(scores, dim=-1)
return attn @ V, attn
That scale factor √d_k is the difference between a model that trains and a model that quietly dies at initialization.
Multi-head, but plain-english
Instead of one big lookup, we do h smaller lookups in parallel. Each head learns a different "aspect" of relevance — one head might track syntax, another might track long-range topics.
class MultiHeadAttention(torch.nn.Module):
def __init__(self, d_model, n_heads):
super().__init__()
assert d_model % n_heads == 0
self.d_head = d_model // n_heads
self.n_heads = n_heads
self.qkv = torch.nn.Linear(d_model, 3 * d_model)
self.out = torch.nn.Linear(d_model, d_model)
def forward(self, x, mask=None):
B, T, C = x.shape
qkv = self.qkv(x).reshape(B, T, 3, self.n_heads, self.d_head)
q, k, v = qkv.permute(2, 0, 3, 1, 4)
out, _ = attention(q, k, v, mask)
out = out.transpose(1, 2).reshape(B, T, C)
return self.out(out)
Positional encoding
Attention alone is order-blind. We inject positions — either fixed sinusoids or learned embeddings. Learned embeddings tend to win in practice, but sinusoids are elegant.
The rest is plumbing
- Feed-forward: two-layer MLP applied per-token.
- Residuals + LayerNorm: keeps gradients healthy.
- Masking: causal for LMs, padding masks for batching.
What actually matters when training
- Weight tying between embeddings and the output head.
- Warmup + cosine decay for the learning rate.
- Gradient clipping at 1.0.
- Mixed precision — free 2x speedup on modern GPUs.
TL;DR
If you can explain scaled dot-product attention out loud, you understand ~70% of Transformers. The rest is polish and vibes.
Happy training ✨