We all use ChatGPT, Gemini, Claude — these large language models that feel almost sentient. You give them a prompt, and they generate coherent, context-aware text one token at a time. It feels like magic.
But have you ever wondered what's actually happening inside that model? Not the API call. Not the wrapper. The actual architecture — the matrix multiplications, the attention mechanisms, the gradient updates that turn random noise into Shakespeare?
I wanted to understand it at the deepest level. So I built a decoder-only Transformer language model from scratch in PyTorch. No nn.TransformerDecoder. No Hugging Face. Every single component — from the embedding layers to the causal self-attention mechanism to the layer normalization — implemented from first principles.
I call it Mini-GPT. It has ~10.7 million parameters, trains on the Tiny Shakespeare dataset, and generates surprisingly coherent Shakespearean text. In this post, I'm going to walk you through every single layer — from input tokens to generated output.
01 —1. The Architecture — A Bird's-Eye View
Mini-GPT is a decoder-only Transformer, the same architecture family that powers GPT-2, GPT-3, and GPT-4. The "decoder-only" part means it only uses the right half of the original Transformer from the "Attention Is All You Need" paper — the part responsible for generating sequences, not encoding them.
Here's the full architecture, top to bottom:
The flow is deceptively simple: tokens go in, logits come out. But the devil is in the six Transformer blocks in the middle — each one containing a multi-head attention mechanism and a feed-forward network, both wrapped in layer normalization and residual connections.
02 —2. Key Design Decisions
Before writing a single line of code, I had to make several architectural choices. Each one has a direct, measurable impact on training stability, convergence speed, and the quality of generated text.
| Decision | Choice | Why |
|---|---|---|
| Norm placement | Pre-LN (GPT-2 style) | Normalizing before each sub-layer produces more stable gradients than Post-LN. The original Transformer used Post-LN, but GPT-2 switched to Pre-LN because it eliminates the need for careful learning rate warmup. |
| Activation | GELU | GELU (Gaussian Error Linear Unit) is smoother than ReLU at the origin. Unlike ReLU which has a hard zero cutoff, GELU smoothly approaches zero, allowing small negative gradients to flow. Standard for all GPT variants. |
| Positional encoding | Learned embeddings | The original Transformer used sinusoidal functions. GPT-2 replaced them with learned position embeddings — a simple lookup table of (position → vector). Simpler and equally effective at this scale. |
| Weight tying | Embedding ↔ LM Head | The token embedding matrix and the final linear head share the same weight matrix. This reduces total parameters by ~25K and forces the model to learn embeddings that are directly useful for prediction. |
| QKV projection | Fused (single linear) | Instead of three separate nn.Linear layers for Q, K, and V, I use one Linear(384, 1152) and then .chunk(3). One GEMM instead of three — significantly more GPU-friendly. |
| Tokenization | Character-level | With only ~1MB of training data, a BPE tokenizer would create a huge vocabulary with sparse statistics. Character-level encoding gives us a tight vocab of 65 unique characters — every character gets rich training signal. |
03 —3. Building Every Component from Scratch
This is the part that makes Mini-GPT different from a tutorial project. Every module in the model/ directory is handwritten — no PyTorch built-in Transformer layers anywhere.
3.1 — Token Embedding
The very first step is converting raw integer token IDs into dense vectors. The implementation is a thin wrapper around nn.Embedding — a lookup table where row i contains the 384-dimensional learned vector for token i.
class Embedding(nn.Module):
def __init__(self, vocab_size: int, embedding_dim: int):
super().__init__()
self.embedding = nn.Embedding(vocab_size, embedding_dim)
def forward(self, token_ids):
return self.embedding(token_ids) # (B, S) → (B, S, 384)
With 65 characters and 384 dimensions, this creates a 65×384 = 24,960 parameter matrix. Small, but critical — these are the vectors the model learns to associate with each character.
3.2 — Positional Embedding
Transformers have no inherent notion of position. Without positional information, the model would treat "ROMEO" and "OEMOR" identically. The solution: a second embedding table indexed by position rather than token.
class PositionalEmbedding(nn.Module):
def __init__(self, max_seq_length: int, embedding_dim: int):
super().__init__()
self.embedding = nn.Embedding(max_seq_length, embedding_dim)
def forward(self, sequence_length: int):
positions = torch.arange(sequence_length, device=self.embedding.weight.device)
return self.embedding(positions) # (S,) → (S, 384)
Position 0 gets one learned vector, position 1 gets another, up to position 255. These are added element-wise to the token embeddings, giving the model a sense of where each token sits in the sequence.
3.3 — Layer Normalization (from scratch)
This is where I drew the line and said "no shortcuts." Instead of using PyTorch's nn.LayerNorm, I implemented it manually:
class LayerNorm(nn.Module):
def __init__(self, embedding_dim: int, eps: float = 1e-5):
super().__init__()
self.eps = eps
self.gamma = nn.Parameter(torch.ones(embedding_dim)) # scale
self.beta = nn.Parameter(torch.zeros(embedding_dim)) # shift
def forward(self, x):
mean = x.mean(dim=-1, keepdim=True)
variance = x.var(dim=-1, unbiased=False, keepdim=True)
normalized = (x - mean) / torch.sqrt(variance + self.eps)
return self.gamma * normalized + self.beta
The math is straightforward: subtract the mean, divide by the standard deviation (with an epsilon to prevent division by zero), then scale and shift with learned gamma and beta parameters. This stabilizes training by ensuring each layer's inputs have consistent statistics.
3.4 — Multi-Head Causal Self-Attention
This is the heart of the Transformer. This is where the model decides which tokens to pay attention to when processing each position.
The implementation is the most complex component. Here's what happens step by step:
- Fused QKV projection: A single
Linear(384, 1152)produces Q, K, and V in one matrix multiply, then.chunk(3)splits them. - Reshape into heads: Each of Q, K, V is reshaped from
(B, S, 384)→(B, 6, S, 64). Each of the 6 heads operates on a 64-dimensional slice. - Scaled dot-product attention:
Q @ K.transpose(-2, -1) / sqrt(64)produces raw attention scores. The scaling prevents the dot products from becoming too large, which would push softmax into regions with tiny gradients. - Causal masking: A pre-registered upper-triangular boolean mask sets future positions to
-inf, ensuring the model can only attend to tokens it has already seen. This is what makes it autoregressive. - Softmax + dropout: Attention weights are normalized and regularized.
- Weighted sum of values:
attention_weights @ Vproduces the output for each head. - Concatenate and project: The 6 heads are concatenated back to 384 dimensions and passed through a final linear projection.
# The critical lines from forward():
scores = Q @ K.transpose(-2, -1) / math.sqrt(self.head_dim)
scores = scores.masked_fill(self.causal_mask[:seq_len, :seq_len], float("-inf"))
attention_weights = torch.softmax(scores, dim=-1)
output = attention_weights @ V
The causal mask is the key difference between a decoder Transformer and an encoder Transformer. In BERT (an encoder), every token can attend to every other token. In GPT (a decoder), token i can only attend to tokens 0 through i. This forces the model to learn to predict the future based only on the past — which is exactly what language modeling requires.
3.5 — Feed-Forward Network
After attention, each position passes through a simple two-layer neural network with a 4× expansion factor:
class FeedForward(nn.Module):
def __init__(self, embedding_dim: int, expansion_factor: int = 4, dropout: float = 0.0):
super().__init__()
hidden_dim = embedding_dim * expansion_factor # 384 → 1536
self.net = nn.Sequential(
nn.Linear(embedding_dim, hidden_dim), # project up
nn.GELU(), # non-linearity
nn.Linear(hidden_dim, embedding_dim), # project back down
nn.Dropout(dropout),
)
The expansion from 384 → 1536 → 384 gives each position a much larger representational space to "think" in, before compressing back down. This is where the model stores most of its learned knowledge.
3.6 — The Transformer Block
A single block wires everything together with Pre-LN residual connections:
class TransformerBlock(nn.Module):
def forward(self, x):
# Multi-Head Attention with Pre-LN residual
x = x + self.attention(self.norm1(x))
# Feed-Forward with Pre-LN residual
x = x + self.ffn(self.norm2(x))
return x
Those two lines are the entire forward pass. The residual connections (x + ...) are critical — they allow gradients to flow directly through the network without vanishing across 6 stacked blocks.
04 —4. The Training Pipeline
Building the model is half the battle. Training it is the other half.
Data preparation
The Tiny Shakespeare dataset is a single 1.1MB text file containing all of Shakespeare's plays. The tokenizer simply maps each unique character to an integer (A=0, B=1, ..., z=64), giving us a vocabulary of 65 tokens.
The data is split 90/10 into training and validation. During training, random contiguous chunks of 256 characters are sampled as input, with the target being the same chunk shifted right by one position:
input = [t₀, t₁, t₂, ..., t₂₅₅] # "ROMEO: Wh..."
target = [t₁, t₂, t₃, ..., t₂₅₆] # "OMEO: Wha..."
Optimizer and schedule
AdamW with weight decay of 0.1 — the standard choice for Transformer training. The learning rate follows a cosine decay schedule with linear warmup: it ramps linearly from 0 to 3×10⁻⁴ during warmup, then cosine-decays to a minimum of 3×10⁻⁵ over the remaining iterations.
Why cosine decay? Because a fixed learning rate either converges too fast (high LR → instability) or too slow (low LR → never reaches minimum). Cosine decay gives you the best of both worlds: aggressive exploration early, then fine-grained convergence.
Training results
| Metric | Start | End | Notes |
|---|---|---|---|
| Train loss | ~4.17 | ~1.4 | Initial = ln(65) = random guessing |
| Val loss | ~4.17 | ~1.5 | Slight gap → mild overfitting (expected) |
| Parameters | 10,742,784 | ~10.7M trainable | |
| Training time (GPU) | ~3-5 minutes | On NVIDIA GPU with CUDA | |
| Training time (CPU) | ~30-45 minutes | On modern CPU (Apple M-series) | |
The starting loss of ~4.17 makes perfect mathematical sense: with a vocabulary of 65 characters, random guessing gives you -ln(1/65) ≈ 4.17. Watching the loss drop from 4.17 to 1.4 means the model went from random noise to confidently predicting the next character with about 75% probability on average.
05 —5. Generation — One Token at a Time
Once trained, the model generates text autoregressively — one character at a time, feeding each generated character back as input for the next prediction.
The generation function supports two critical controls:
- Temperature — A scalar that divides the logits before softmax. Temperature = 1.0 gives you the raw learned distribution. Temperature = 0.5 sharpens it (more deterministic, "safer" text). Temperature = 1.5 flattens it (more creative, riskier).
- Top-K — Zeros out all logits below the K-th highest value. Top-K = 20 means only the 20 most probable characters are candidates. This prevents the occasional completely random character that ruins coherence.
# Temperature scaling + Top-K filtering
logits = logits / temperature
if top_k is not None:
top_k_values, _ = torch.topk(logits, min(top_k, logits.size(-1)))
threshold = top_k_values[:, -1].unsqueeze(-1)
logits[logits < threshold] = float("-inf")
probs = F.softmax(logits, dim=-1)
next_token = torch.multinomial(probs, num_samples=1)
The magic moment? When the model — having only seen raw characters and their statistical relationships — starts outputting correctly formatted Shakespearean plays. Character names in uppercase. Proper line breaks. Archaic English vocabulary. Appropriate dramatic dialogue. It learned all of this purely from next-character prediction.
06 —6. What I Learned
Building Mini-GPT from scratch taught me things no tutorial or API wrapper ever could:
- Attention is just matrix multiplication. The mystique around "self-attention" evaporates when you implement it. It's Q×Kᵀ → softmax → ×V. Three matrix multiplies and a normalization.
- Residual connections are non-negotiable. Without the
x + ...skip connections, gradients vanish by the third block. The network literally cannot learn. - Pre-LN vs Post-LN changes everything. Post-LN (the original paper) requires very careful warmup. Pre-LN (GPT-2 style) just works. A one-line change that eliminates hours of hyperparameter tuning.
- Weight tying is free performance. Sharing the embedding matrix with the output head doesn't just save parameters — it actively improves generation quality because the output space is forced to be semantically consistent with the input space.
- The causal mask is what makes it generative. The upper-triangular mask of
-infvalues is the entire difference between BERT and GPT. Without it, you have an encoder. With it, you have a generator.
07 —The Complete Project Structure
mini-gpt/
├── config.py # Model & training configuration
├── train.py # End-to-end training script
├── inference.py # Text generation from checkpoints
├── model/
│ ├── embedding.py # Token embedding (from scratch)
│ ├── positional_embedding.py # Learned position encoding
│ ├── layer_norm.py # LayerNorm (from scratch)
│ ├── multi_head_attention.py # Multi-head causal self-attention
│ ├── feed_forward.py # Position-wise FFN (GELU)
│ ├── transformer_block.py # Pre-LN Transformer block
│ └── gpt.py # Full GPT model + generation
├── tokenizer/
│ ├── tokenizer.py # Character-level tokenizer
│ └── bpe.py # BPE exploration
├── training/
│ └── dataset.py # Data loading & batching
├── tests/
│ ├── test_model.py # Architecture tests
│ └── test_tokenizer.py # Tokenizer tests
└── data/
└── input.txt # Tiny Shakespeare (~1MB)
Building Mini-GPT was one of the most rewarding engineering exercises I've done. When you strip away the APIs, the frameworks, the managed services — and build a language model from raw matrix operations — you realize something profound. There is no magic. There is no intelligence. There is only statistics, linear algebra, and gradient descent. And somehow, from those three ingredients, something that looks remarkably like understanding emerges.
That's the real lesson of building from scratch. Not the code. The understanding.
Let's build smart. Let's build together.
— Gopal