Build a Transformer from Scratch: The TinyLLM Guide

Why Train a Transformer on Mathematical Equations?

Most LLM tutorials start by throwing massive novels or scrapings of Shakespeare at a neural network. This makes it impossible for beginners to spot when the model actually learns sequence order or simply regurgitates strings.

With **TinyLLM**, we scale the problem down to its mathematical bedrock. We train a tiny model (~10,000 parameters) to solve simple equations like 2 + 2 = 4 and 2 + 3 = 5. Because the vocabulary is small (only 17 characters) and the mathematical rules are rigid, we can watch the cross-entropy loss collapse from random noise to 0.000 in real time.

The Transformer Block Architecture Flow

A standard modern GPT-style Transformer block stacks self-attention layers with residual connections, layer normalization (specifically `RMSNorm`), and a feed-forward multi-layer perceptron (MLP). Below is the logical data flow of the TinyLLM block:

The 10-Step Checklist: Coding the Transformer

Let's walk through the exact coding process of building `TinyLLM`. All code snippets are extracted directly from the working files in our codebase.

# Vocabulary tokens
<pad>
<bos>
<eos>
2
3
+
=
4
5

# train_v2.py
vocab = open("vocab.txt").read().splitlines()
stoi = {s:i for i,s in enumerate(vocab)}
itos = {i:s for i,s in enumerate(vocab)}

# Mapping a math equation:
# " 2 + 2 = 4 " -> [1, 9, 11, 9, 12, 13, 2]

# model_v2.py
# Token Embedding
self.embed = nn.Embedding(self.vocab_size, self.hidden_size)
# Position Embedding
self.pos_embed = nn.Embedding(max_seq_len, self.hidden_size)

# Forward pass combining them:
pos = torch.arange(t, device=x.device)
x = self.embed(x) + self.pos_embed(pos)

# Projection definitions
self.q_proj = nn.Linear(8, 8, bias=False)
self.k_proj = nn.Linear(8, 8, bias=False)
self.v_proj = nn.Linear(8, 8, bias=False)

# Computation
Q = self.q_proj(x)
K = self.k_proj(x)
V = self.v_proj(x)

# Reshape for multihead: shape (Batch, Seq, Heads, HeadDim)
B, T, C = Q.shape
Q = Q.view(B, T, 2, 4).transpose(1, 2)
K = K.view(B, T, 2, 4).transpose(1, 2)
V = V.view(B, T, 2, 4).transpose(1, 2)

# Scaled dot-product
scores = (Q @ K.transpose(-2, -1)) / math.sqrt(4)

# Force future elements to -infinity
mask = torch.tril(torch.ones(T, T)).to(x.device)
scores = scores.masked_fill(mask == 0, -1e9)

attn = torch.softmax(scores, dim=-1)
out = attn @ V

# Concat heads back into hidden size channel
out = out.transpose(1, 2).contiguous().view(B, T, C)
out = self.o_proj(out)

# Forward norm step
x = residual + out

# RMSNorm Equation
class RMSNorm(nn.Module):
    def forward(self, x):
        rms = x.pow(2).mean(-1, keepdim=True).sqrt()
        return x / (rms + self.eps) * self.scale

residual = x
x = self.norm2(x)
x = self.mlp_up(x)
x = torch.relu(x)
x = self.mlp_down(x)
x = residual + x

# train_v2.py training step
logits = model(X) # shape: (B, SeqLen, VocabSize)

# Flatten sequence length and batches together for loss evaluation
loss = loss_fn(logits.view(-1, len(vocab)), Y.view(-1))
optimizer.zero_grad()
loss.backward()
optimizer.step()

Reference Code: The Complete Runnable Python Model

Below is the complete, self-contained Python script implementing the TinyLLM architecture, dataset preparation, and training loop. You can copy-paste and run this locally in your terminal or a Jupyter notebook. It has zero external dataset dependencies and trains in under a second on standard CPUs:

import torch
import torch.nn as nn
import torch.optim as optim
import math

# 1. Define Vocabulary and Training Corpus in-memory
vocab = ["", "", "", "hello", "world", "I", "am", "tiny", "model", "2", "3", "+", "=", "4", "5", "?", "."]
stoi = {s: i for i, s in enumerate(vocab)}
itos = {i: s for i, s in enumerate(vocab)}

# Mathematical training corpus
corpus = [
    " hello world ",
    " I am tiny model ",
    " 2 + 2 = 4 ",
    " 2 + 2 = 4 ",
    " 2 + 3 = 5 ",
    " 2 + 3 = 5 "
]

# 2. Tokenize corpus
data = []
for line in corpus:
    tokens = line.split()
    ids = [stoi[t] for t in tokens if t in stoi]
    data.append(ids)

# 3. Create Training Pairs (X: Query Sequence, Y: Target Sequence)
X_list, Y_list = [], []
for seq in data:
    if len(seq) < 2: continue
    X_list.append(seq[:-1]) # Input: " 2 + 2 ="
    Y_list.append(seq[1:])  # Target: "2 + 2 = 4"

# Pad sequences (ignore padding in loss function using -100)
max_len = max(len(x) for x in X_list)
X_padded = [x + [stoi[""]] * (max_len - len(x)) for x in X_list]
Y_padded = [y + [-100] * (max_len - len(y)) for y in Y_list]

X = torch.tensor(X_padded)
Y = torch.tensor(Y_padded)

# 4. Define the Model Architecture
class RMSNorm(nn.Module):
    def __init__(self, dim, eps=1e-6):
        super().__init__()
        self.scale = nn.Parameter(torch.ones(dim))
        self.eps = eps

    def forward(self, x):
        rms = x.pow(2).mean(-1, keepdim=True).sqrt()
        return x / (rms + self.eps) * self.scale

class TinyLLM(nn.Module):
    def __init__(self, vocab_size=17, max_seq_len=64):
        super().__init__()
        self.vocab_size = vocab_size
        self.hidden_size = 8
        self.num_heads = 2
        self.head_dim = 4

        self.embed = nn.Embedding(vocab_size, 8)
        self.pos_embed = nn.Embedding(max_seq_len, 8)

        self.q_proj = nn.Linear(8, 8, bias=False)
        self.k_proj = nn.Linear(8, 8, bias=False)
        self.v_proj = nn.Linear(8, 8, bias=False)
        self.o_proj = nn.Linear(8, 8, bias=False)

        self.norm1 = RMSNorm(8)
        self.norm2 = RMSNorm(8)

        self.mlp_up = nn.Linear(8, 17, bias=False)
        self.mlp_down = nn.Linear(17, 8, bias=False)
        self.lm_head = nn.Linear(8, vocab_size, bias=False)

    def forward(self, x):
        b, t = x.shape
        pos = torch.arange(t, device=x.device)
        x = self.embed(x) + self.pos_embed(pos)

        # Attention block
        residual = x
        x = self.norm1(x)
        Q = self.q_proj(x).view(b, t, 2, 4).transpose(1, 2)
        K = self.k_proj(x).view(b, t, 2, 4).transpose(1, 2)
        V = self.v_proj(x).view(b, t, 2, 4).transpose(1, 2)

        scores = (Q @ K.transpose(-2, -1)) / math.sqrt(4)
        mask = torch.tril(torch.ones(t, t, device=x.device))
        scores = scores.masked_fill(mask == 0, -1e9)

        attn = torch.softmax(scores, dim=-1)
        out = (attn @ V).transpose(1, 2).contiguous().view(b, t, 8)
        x = residual + self.o_proj(out)

        # MLP block
        residual = x
        x = self.norm2(x)
        x = self.mlp_down(torch.relu(self.mlp_up(x)))
        x = residual + x

        return self.lm_head(x)

# 5. Training Loop
model = TinyLLM(vocab_size=len(vocab))
loss_fn = nn.CrossEntropyLoss(ignore_index=-100)
optimizer = optim.Adam(model.parameters(), lr=0.01)

print("Training started...")
for epoch in range(801):
    logits = model(X)
    loss = loss_fn(logits.view(-1, len(vocab)), Y.view(-1))
    
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

    if epoch % 200 == 0:
        print(f"Epoch {epoch:04d} | Loss: {loss.item():.4f}")

# 6. Run Inference
model.eval()
test_tokens = [stoi[""], stoi["2"], stoi["+"], stoi["2"]]
generated = test_tokens.copy()

with torch.no_grad():
    for _ in range(3):
        inp = torch.tensor([generated])
        logits = model(inp)
        next_id = torch.argmax(logits[0, -1, :]).item()
        generated.append(next_id)
        if next_id == stoi[""]: break

print("\nPrompt: ['', '2', '+', '2']")
print("Generated completion:", [itos[idx] for idx in generated])