AI/ML

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.

KeerthanMay 18, 20252 min read
transformerspytorchdeep-learning

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

  1. Weight tying between embeddings and the output head.
  2. Warmup + cosine decay for the learning rate.
  3. Gradient clipping at 1.0.
  4. 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 ✨