TinyGPT — Interactive Anatomy & Playground

Introduction Overview

What this demo actually builds

This app is a tiny decoder-only Transformer trained from scratch on a 10-sentence English corpus. It learns to predict the next token given a context window of 6 tokens. Despite its size (~few thousand parameters), it demonstrates every core idea behind modern LLMs: token embeddings, positional embeddings, causal self-attention, multi-head attention, feed-forward blocks, layer normalization, and autoregressive generation.

📦 Corpus

10 short sentences, ~42 unique words, <END> token appended.

🧱 Block size

6 tokens of context. The model only ever "sees" this window.

🧠 Embedding

32-dim vectors for tokens and positions — summed before attention.

👁️ Attention

2 heads, 2 layers, causal mask so each token only attends to past.

⚡ Training

AdamW, cross-entropy loss, 1500 steps on random batches.

✨ Generation

Autoregressive sampling with softmax + multinomial.

High-level pipeline

Text

→

Tokens

→

Embed + Pos

→

×N Blocks

→

LayerNorm

→

Linear Head

→

Logits

Corpus & Vocabulary Data Prep

Turning sentences into a numbered dictionary

1. The raw corpus

Ten short English sentences, each terminated with a special <END> token so the model knows when a sentence finishes.

# 10 sentences, each appended with <END> corpus = [ "hello friends how are you ", "the coffee is very hot ", ... ] corpus = [s + " <END>" for s in corpus]

2. Build the vocabulary

We split on whitespace, take the set of unique words, and assign each an integer ID. This gives us word2idx (word → int) and idx2word (int → word).

words = list(set(text.split())) vocab_size = len(words) word2idx = {w: i for i, w in enumerate(words)} idx2word = {i: w for w, i in word2idx.items()}

Live vocabulary — click any token to highlight its ID mapping:

Data & Batching Tensors

From words to a 1-D tensor, then to (x, y) pairs

1. Encode the whole corpus

The entire text is flattened into one long 1-D tensor of integer IDs. This is our training "tape".

data = torch.tensor([word2idx[w] for w in text.split()], dtype=torch.long)

2. `get_batch()` — sliding windows

We sample batch_size random starting positions. For each, x is 6 tokens and y is the same window shifted by 1 — the "next token" targets.

def get_batch(batch_size=16): ix = torch.randint(len(data) - block_size, (batch_size,)) x = torch.stack([data[i:i+block_size] for i in ix]) y = torch.stack([data[i+1:i+block_size+1] for i in ix]) return x, y

Each row = one training example. x is the context, y is what the model must predict.

TinyGPT nn.Module

The top-level model: embeddings → blocks → head

Constructor — what it builds

class TinyGPT(nn.Module): def __init__(self): super().__init__() self.token_embedding = nn.Embedding(vocab_size, embedding_dim) self.position_embedding = nn.Embedding(block_size, embedding_dim) self.blocks = nn.Sequential(*[Block(embedding_dim, block_size, n_heads) for _ in range(n_layers)]) self.ln_f = nn.LayerNorm(embedding_dim) self.head = nn.Linear(embedding_dim, vocab_size)

🔤 token_embedding

A lookup table of shape [vocab_size × 32]. Each word ID maps to a 32-dim learned vector.

📍 position_embedding

A lookup table of shape [block_size × 32]. Tells the model "you are token 0, 1, 2…" — because attention itself is permutation-invariant.

🧱 blocks

A stack of 2 Transformer blocks (attention + FFN + residuals + LayerNorm).

🎯 head

A final linear layer that projects each 32-dim hidden state back to vocab_size logits — one score per word.

Hyperparameters

block_size

6 — context window

embedding_dim

32 — hidden width

n_heads

2 — attention heads

n_layers

2 — transformer blocks

lr

1e-3 — AdamW learning rate

epochs

1500 — training steps

SelfAttentionHead Core

The fundamental "token talks to token" operation

What it holds

class SelfAttentionHead(nn.Module): def __init__(self, embedding_dim, block_size, head_size): self.key = nn.Linear(embedding_dim, head_size, bias=False) self.query = nn.Linear(embedding_dim, head_size, bias=False) self.value = nn.Linear(embedding_dim, head_size, bias=False) self.register_buffer('tril', torch.tril(torch.ones(block_size, block_size)))

🔑 Key (K)

"What do I contain?" — each token advertises itself.

❓ Query (Q)

"What am I looking for?" — each token searches.

💎 Value (V)

"What info do I actually pass on if attended to?"

The forward pass — scaled dot-product attention

def forward(self, x): k = self.key(x) q = self.query(x) wei = q @ k.transpose(-2, -1) / (C ** 0.5) # QKᵀ / √d wei = wei.masked_fill(self.tril[:T, :T] == 0, float('-inf')) # causal mask wei = F.softmax(wei, dim=-1) # attention weights v = self.value(x) return wei @ v # weighted sum of values

Interactive: the causal mask

The lower-triangular mask ensures token i can only attend to tokens ≤ i. This is what makes the model a decoder — it can't peek at the future.

Causal mask (tril)

Allowed — token can attend

Masked (−∞) — future, blocked

After softmax, masked cells become exactly 0. The remaining cells in each row sum to 1.

MultiHeadAttention Ensemble

Several attention heads running in parallel, then concatenated

class MultiHeadAttention(nn.Module): def __init__(self, embedding_dim, block_size, num_heads): head_size = embedding_dim // num_heads self.heads = nn.ModuleList([SelfAttentionHead(embedding_dim, block_size, head_size) for _ in range(num_heads)]) self.proj = nn.Linear(num_heads * head_size, embedding_dim) def forward(self, x): out = torch.cat([h(x) for h in self.heads], dim=-1) return self.proj(out)

Why multiple heads?

Each head can specialize — one might track subject-verb agreement, another might track nearby adjectives. With embedding_dim=32 and num_heads=2, each head has head_size=16. Their outputs are concatenated back to 32 dims and projected through a linear layer so the next block sees a clean 32-dim signal.

x (32)

→

Head 1 (16)

↘

concat (32)

→

proj (32)

↗

Head 2 (16)

FeedForward Per-token MLP

A tiny MLP applied independently to every token

class FeedForward(nn.Module): def __init__(self, n_embd): self.net = nn.Sequential( nn.Linear(n_embd, 4 * n_embd), nn.ReLU(), nn.Linear(4 * n_embd, n_embd) ) def forward(self, x): return self.net(x)

Think of attention as "gathering information from other tokens" and the FFN as "thinking about what I gathered". The 4× expansion (32 → 128 → 32) is a standard trick that gives the model extra capacity per token without increasing sequence-level cost much.

With n_embd=32:

Input

32

Hidden

128

Output

32

Transformer Block Composition

Attention + FFN, each wrapped with residual + LayerNorm (Pre-LN)

class Block(nn.Module): def __init__(self, embedding_dim, block_size, n_heads): self.sa = MultiHeadAttention(embedding_dim, block_size, n_heads) self.ffwd = FeedForward(embedding_dim) self.ln1 = nn.LayerNorm(embedding_dim) self.ln2 = nn.LayerNorm(embedding_dim) def forward(self, x): x = x + self.sa(self.ln1(x)) # residual around attention x = x + self.ffwd(self.ln2(x)) # residual around FFN return x

Why residuals?

Each sub-layer gets a "skip connection": x + sublayer(x). This lets gradients flow directly through the network and lets the model learn small refinements on top of the identity — critical for training deep stacks.

Why Pre-LN (LayerNorm before the sublayer)?

Modern GPTs normalize inputs to each sublayer rather than outputs. It stabilizes training and is the dominant convention since GPT-2.

x

→

LN₁

→

MultiHeadAttn

→

+ x (residual)

→

LN₂

→

FFN

→

+ x (residual)

forward() Inference

One complete pass from token IDs to logits + loss

def forward(self, idx, targets=None): B, T = idx.shape tok_emb = self.token_embedding(idx) pos_emb = self.position_embedding(torch.arange(T, device=idx.device)) x = tok_emb + pos_emb x = self.blocks(x) x = self.ln_f(x) logits = self.head(x) loss = None if targets is not None: B, T, C = logits.shape loss = F.cross_entropy(logits.view(B*T, C), targets.view(B*T)) return logits, loss

Step by step

1. Embed tokens

Look up each of the B×T token IDs → shape [B, T, 32].

2. Add position

Broadcast the [T, 32] positional embedding across the batch.

3. Run blocks

Sequentially through 2 transformer blocks.

4. Final LayerNorm

Stabilize the hidden states before the head.

5. Linear head

Project 32 → vocab_size. Output shape [B, T, vocab_size].

6. Cross-entropy loss

Flatten to [B*T, vocab_size] vs [B*T] targets.

generate() Autoregressive

Sampling one token at a time, feeding it back in

def generate(self, idx, max_new_tokens): for _ in range(max_new_tokens): idx_cond = idx[:, -block_size:] # crop to context window logits, _ = self(idx_cond) logits = logits[:, -1, :] # only the last position probs = F.softmax(logits, dim=-1) next_idx = torch.multinomial(probs, 1) # sample idx = torch.cat((idx, next_idx), dim=1) # append return idx

Starting from a seed word, the model predicts the next word, appends it, re-runs on the last 6 tokens, predicts again, and so on. Because we sample (not argmax), outputs vary between runs.

Training Loop Optimization

AdamW, cross-entropy, 1500 steps

optimizer = torch.optim.AdamW(model.parameters(), lr=lr) for step in range(epochs): xb, yb = get_batch() logits, loss = model(xb, yb) optimizer.zero_grad() loss.backward() optimizer.step() if step % 300 == 0: print(f"Step {step}, loss={loss.item():.4f}")

Simulated loss curve

Typical shape: starts high (≈ vocab-level random), drops quickly, then slowly refines.

Interactive Generation Live Simulation

Enter any of the 42 vocabulary words and watch the LLM simulate autoregressive generation step-by-step, complete with probability distributions.