Mini GPT (Character-Level Language Model) - End-to-End Implementations

Implement Mini GPT (Character-Level Language Model) in Python with NumPy. hard exercise with examples, hints, and instant feedback. No setup or signup required.

Mini GPT: a Character-Level Language Model

What you'll build

A tiny decoder-only Transformer with the same architecture as GPT-2 (just very small), run end to end: forward pass, next-token loss, a training loop and greedy generation.

Architecture

Step	Operation	Shape
input	`idx`: integer token ids, `T ≤ block_size`	`(B, T)`
embeddings	`x = tok_emb[idx] + pos_emb[:T]`	`(B, T, C)`
per block (pre-norm)	`x = x + attention(LN1(x))`	`(B, T, C)`
per block (pre-norm)	`x = x + gelu(LN2(x) @ W_1 + b_1) @ W_2 + b_2`	hidden `(B, T, 4C)`
final norm	`h = LN_f(x)` = `hidden_states(idx)`	`(B, T, C)`
LM head (untied)	`logits = h @ W_lm`	`(B, T, V)`

Causal self-attention (per block, `H = n_heads`, `d_head = C / H`)

q, k, v = split(x @ W_qkv, 3, axis=-1): each (B, T, C); the columns of W_qkv are [Q | K | V].
Reshape each to (B, H, T, d_head); head h uses columns h·d_head : (h+1)·d_head.
scores = q @ kᵀ / sqrt(d_head), set to −inf where the key position is after the query position.
out = softmax(scores) @ v, merged back to (B, T, C), then @ W_o. No biases in attention.

The constructor (given) creates every parameter with GPT-2-style N(0, 0.02²) weights from np.random.default_rng(seed); LayerNorm (eps = 1e-5), tanh-GELU and softmax helpers are provided.

Method contract

Method	Returns
`attention(x, blk)`	`(B, T, C)`, using the block dict's `W_qkv` and `W_o`
`hidden_states(idx)`	`(B, T, C)`: embeddings → blocks → final LayerNorm
`forward(idx)`	logits `(B, T, vocab_size)`
`loss(idx, targets)`	float: mean cross-entropy over all `B·T` positions; `targets[b, t]` is the token after `idx[b, t]`
`generate(idx, max_new_tokens)`	int array `(B, T + max_new_tokens)` starting with `idx`
`train_step(idx, targets, lr)`	(provided in the NumPy track) one SGD step on `W_lm` only, using `dL/dlogits = (softmax − one_hot) / N`
`train(model, idx, targets, steps=20, lr=1.0)`	list of `steps` losses, each computed before that step's update

Conventions

Initial loss ≈ log(V): with std-0.02 weights the logits start near 0, so the model guesses uniformly.
Greedy generation: repeat max_new_tokens times: crop the context to the last block_size tokens, run forward, append the argmax of the last position's logits.
Use a numerically stable log-softmax in loss (subtract the max).

Constraints

NumPy only; no loops over batch, heads or positions (only over blocks and generation steps).

Follow-up

Generation re-runs the whole context each step: how does a KV cache make each new token O(T) instead of O(T²)?
Why do GPT-2 and most LLMs tie the LM head to tok_emb? (See Transformer Bug Hunt, which does.)

Examples

Input: model = MiniGPT(vocab_size=5); model.forward(np.zeros((2, 8), dtype=int)).shape
Output: (2, 8, 5)

One distribution over the 5 tokens for every position of every sequence

Input: seq = np.array([[i % 5 for i in range(9)]]); model.loss(seq[:, :-1], seq[:, 1:])
Output: 1.6769

Close to log 5 = 1.6094: the untrained model is nearly uniform over the 5 tokens

Input: After train(model, seq[:, :-1], seq[:, 1:], steps=30, lr=2.0) on two shifted copies of the cycle 0→1→2→3→4→0: model.generate(np.array([[0, 1, 2]]), 5)
Output: [[0, 1, 2, 3, 4, 0, 1, 2]]

The prompt is kept and greedy decoding continues the learned cycle

Practice in Python

import numpy as np

def layer_norm(x, gamma, beta, eps=1e-5):
    mu = x.mean(axis=-1, keepdims=True)
    var = x.var(axis=-1, keepdims=True)
    return (x - mu) / np.sqrt(var + eps) * gamma + beta


def gelu(x):
    # tanh approximation (GPT-2)
    return 0.5 * x * (1.0 + np.tanh(np.sqrt(2.0 / np.pi) * (x + 0.044715 * x ** 3)))


def softmax(x, axis=-1):
    e = np.exp(x - x.max(axis=axis, keepdims=True))
    return e / e.sum(axis=axis, keepdims=True)


class MiniGPT:
    """A tiny character-level GPT: pre-norm decoder blocks, untied LM head."""

    def __init__(self, vocab_size, d_model=16, n_heads=2, n_layers=1, block_size=8, seed=0):
        rng = np.random.default_rng(seed)

        def normal(*shape):
            return rng.normal(0.0, 0.02, size=shape)  # GPT-2 style small init

        self.vocab_size, self.d_model = vocab_size, d_model
        self.n_heads, self.block_size = n_heads, block_size
        self.tok_emb = normal(vocab_size, d_model)   # (V, C)
        self.pos_emb = normal(block_size, d_model)   # (block_size, C)
        self.blocks = []
        for _ in range(n_layers):
            self.blocks.append({
                'ln1_g': np.ones(d_model), 'ln1_b': np.zeros(d_model),
                'W_qkv': normal(d_model, 3 * d_model),    # columns: [Q | K | V]
                'W_o': normal(d_model, d_model),
                'ln2_g': np.ones(d_model), 'ln2_b': np.zeros(d_model),
                'W_1': normal(d_model, 4 * d_model), 'b_1': np.zeros(4 * d_model),
                'W_2': normal(4 * d_model, d_model), 'b_2': np.zeros(d_model),
            })
        self.lnf_g, self.lnf_b = np.ones(d_model), np.zeros(d_model)
        self.W_lm = normal(d_model, vocab_size)      # (C, V)

    def attention(self, x, blk):
        """Causal multi-head self-attention. x: (B, T, C) -> (B, T, C)"""
        # Your code here
        pass

    def hidden_states(self, idx):
        """Embeddings -> blocks -> final LayerNorm. idx: (B, T) -> (B, T, C)"""
        # Your code here
        pass

    def forward(self, idx):
        """Return logits (B, T, vocab_size)."""
        # Your code here
        pass

    def loss(self, idx, targets):
        """Mean next-token cross-entropy (float)."""
        # Your code here
        pass

    def generate(self, idx, max_new_tokens):
        """Greedy decoding. Returns (B, T + max_new_tokens)."""
        # Your code here
        pass

    def train_step(self, idx, targets, lr=1.0):
        """(Provided) One SGD step on the LM head W_lm only; everything else stays frozen.
        For mean cross-entropy, dL/dlogits = (softmax(logits) - one_hot(targets)) / N."""
        h = self.hidden_states(idx)                  # (B, T, C)
        probs = softmax(h @ self.W_lm)               # (B, T, V)
        dlogits = (probs - np.eye(self.vocab_size)[targets]) / targets.size
        grad = h.reshape(-1, self.d_model).T @ dlogits.reshape(-1, self.vocab_size)
        self.W_lm -= lr * grad


def train(model, idx, targets, steps=20, lr=1.0):
    """Call model.train_step `steps` times. Return the list of losses,
    each computed before that step's update."""
    # Your code here
    pass

The interactive editor loads with Python and NumPy. Run your implementation against 9 tests, or use the progressive hints.

More exercises in this section →