Deconstructing TinyGPT from Scratch

Multi-Head Causal Self-Attention

class MultiHeadCausalSelfAttention(nn.Module):
    def __init__(self, n_embd, n_head, block_size, dropout=0.0):
        super().__init__()
        self.n_head = n_head
        self.d_head = n_embd // n_head
        self.qkv = nn.Linear(n_embd, 3 * n_embd)
        self.proj = nn.Linear(n_embd, n_embd)
        mask = torch.tril(torch.ones(block_size, block_size))
        self.register_buffer("mask", mask.view(1, 1, block_size, block_size))

    def forward(self, x):
        B, T, C = x.shape
        q, k, v = self.qkv(x).split(C, dim=2)
        q = q.view(B, T, self.n_head, self.d_head).transpose(1, 2)
        k = k.view(B, T, self.n_head, self.d_head).transpose(1, 2)
        v = v.view(B, T, self.n_head, self.d_head).transpose(1, 2)
        att = (q @ k.transpose(-2, -1)) / math.sqrt(self.d_head)
        att = att.masked_fill(self.mask[:, :, :T, :T] == 0, float("-inf"))
        att = F.softmax(att, dim=-1)
        y = (att @ v).transpose(1, 2).contiguous().view(B, T, C)
        return self.proj(y)

A single linear projects to $[Q \mid K \mid V]$ at once — saves three calls to the dispatcher. The mask is registered as a buffer so it moves with .to(device) but is not a learnable parameter.

The Block

class Block(nn.Module):
    def __init__(self, cfg):
        super().__init__()
        self.ln1 = nn.LayerNorm(cfg.n_embd)
        self.attn = MultiHeadCausalSelfAttention(cfg.n_embd, cfg.n_head,
                                                  cfg.block_size, cfg.dropout)
        self.ln2 = nn.LayerNorm(cfg.n_embd)
        self.mlp = nn.Sequential(
            nn.Linear(cfg.n_embd, 4 * cfg.n_embd), nn.GELU(),
            nn.Linear(4 * cfg.n_embd, cfg.n_embd), nn.Dropout(cfg.dropout))

    def forward(self, x):
        x = x + self.attn(self.ln1(x))
        x = x + self.mlp(self.ln2(x))
        return x

Pre-norm. The $4\times$ MLP expansion is the GPT-2 default. GELU is used because it is what GPT-2 used.

The Model

class TinyGPT(nn.Module):
    def __init__(self, cfg):
        super().__init__()
        self.cfg = cfg
        self.tok_emb = nn.Embedding(cfg.vocab_size, cfg.n_embd)
        self.pos_emb = nn.Embedding(cfg.block_size, cfg.n_embd)
        self.drop = nn.Dropout(cfg.dropout)
        self.blocks = nn.ModuleList([Block(cfg) for _ in range(cfg.n_layer)])
        self.ln_f = nn.LayerNorm(cfg.n_embd)
        self.lm_head = nn.Linear(cfg.n_embd, cfg.vocab_size, bias=False)
        self.lm_head.weight = self.tok_emb.weight   # weight tying

Total parameter count (vocab 12, block 12, 4 layers, 4 heads, $n_\text{embd}=64$): 200,832. Weight tying means the embedding matrix and the LM head share their $12 \times 64$ weights.

Synthetic Dataset

Each example is a 12-character string: sample $A, B \sim \text{Uniform}\{0, \ldots, 999\}$, format as ABC+DEF=, compute $C = A + B$, zero-pad to four digits, reverse. Concatenate to get a 12-character training string.

Supervision masking. We supervise as a standard causal LM but set the target tokens for the prompt portion to $-100$ (CrossEntropyLoss's ignore_index). Only the four answer positions contribute to the loss. The model still attends to the prompt; we just do not train it to reconstruct the prompt.

Training Loop

opt = torch.optim.AdamW(model.parameters(), lr=5e-4, weight_decay=0.01)
sched = torch.optim.lr_scheduler.CosineAnnealingLR(opt, T_max=epochs)
for epoch in range(epochs):
    for x_batch, y_batch in batches:
        _, loss = model(x_batch, y_batch)
        opt.zero_grad(set_to_none=True)
        loss.backward()
        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
        opt.step()
    sched.step()

Standard recipe: AdamW with weight decay $0.01$, cosine annealing, gradient clipping at $1.0$. Nothing exotic. Batch size $128$, $50$ epochs, $10{,}000$ training examples.

Summary

• Attention module: $\sim$30 lines.
• Block: $\sim$15 lines.
• Full model with generation and inspection: $\sim$120 lines.
• Training loop: $\sim$30 lines.

Part 3 runs the experiment, reports the per-position learning curriculum, and inspects the trained attention.

Full code on GitHub: github.com/soveshmohapatra/TinyGPT