Everyone knows what GPT-2 is. We’ll try to understand it by implementing a tiny version of it. The reason I say tiny version is because we can’t really replicate the exact architecture because of limited resources. I only have a single T4 GPU on my Colab and probably only 3 - 4 hours of runtime. So, we’ll use a small dataset and small architecture configuration. I guess we should call it tiny GPT-2.

The goal of this ‘blog’ is not to replicate GPT-2’s full scale. Instead, the goal is to deeply understand:

• What each step in the pipeline does
• Why it is needed
• How it connects to the previous step
• What math happens inside it
• How the code implements that math

If you want to skip the blog and go to the direct implementation, you can find it on my GitHub and on Colab.

I will try to explain every step thoroughly and we will carry one running example throughout:

“I think therefore I am.”

Notation Table

1. Tokenization

1.1 Why Tokenization Exists

Transformers do not understand letters, words, or sentences. They only understand numbers. Tokenization does exactly that:

Convert text into integers in a meaningful, consistent way.

Emphasis on “meaningful”. If “cat” is [20, 3, 99] today and [4, 77] tomorrow, the model cannot learn anything consistent. So, tokenization must be deterministic, reversible and stable across training and inference.

GPT-2 uses byte-level BPE (Byte Pair Encoding), which is efficient and capable of handling all languages, including emojis.

1.2 Byte-Level BPE: Characters → Bytes → Merges

Byte-level BPE works by first turning text into raw bytes (0–255). Then it looks for byte pairs that appear together frequency. When it finds such a pair, it merges those two bytes into a single, larger token and assigns that new token its own integer ID. As this process repeats, the tokenizer gradually builds a small, efficient set of subword tokens.

This approach avoids the problems of character-level tokenization (too long) and word-level tokenization (too many unknown words). Because it starts from bytes, it can represent any text and because it learns common subword patterns, it keeps sequences short. Frequent patterns like “th” or “ing” become single tokens, while rare patterns stay broken into smaller pieces, so the tokenizer never produces an “unknown” token.

Let’s use UTF-8 for encoding:

This is the raw material for BPE.

Let the initial sequence be \(X = [b_1, b_2, \ldots, b_n]\), where \(b_i \in \{0, \ldots, 255\}\) and \(V = \{0, 1, \ldots, 255\}\) be the initial vocabulary of bytes.

For any adjacent pair of bytes \(b_i\) and \(b_{i+1}\), we count the number of times they appear together in the sequence. This is denoted by \(C(b_i, b_{i+1})\).

During each merge, we select the most frequent pair of bytes and merge them into a new token \(t_{new}\) :

and replace every non-overlapping occurrence of \(uv\) with \(t_{new}\) inside \(X\). After merging, the length of \(X\) typically decreases (it “compresses”), and we recompute pair counts on the new sequence. We repeat this until the vocabulary reaches the target size (e.g., 4096).

Intuitively, high-frequency neighboring bytes (or already-merged tokens) become single tokens. Over many merges, common character bigrams, trigrams, and word pieces are “absorbed” into compact subword tokens. This improves efficiency downstream: sequences become shorter, and the model can learn reusable subword patterns.

1.3 Counting Adjacent Pairs

BPE is a greedy algorithm, at each step we only merge the most frequent adjacent pair. This requires recomputing pair statistics after every merge because merges change the sequence and therefore the counts. Selecting the most frequent pair ensures we prioritize the merges that yield the biggest compression and the most reusable subword units.

Example counts:

The most frequent pair is merged into a new token. After a merge:

73, 32 → 256
32, 116 → 257
116, 104 → 258
104, 101 → 259

So the sequence compresses.

This repeats until we reach vocabulary size (1024).

1.4 Why Byte-Level BPE Works

Why not just tokenize by words? or why not just keep characters?

Well, Byte-level BPE answers both concerns. Characters keep sequences too long; words produce out-of-vocabulary problems and brittle vocabularies. Byte-level BPE starts from the universal building blocks (bytes) and learns useful subwords, guaranteeing reversibility and zero unknowns, while keeping sequences reasonably short, almost like a child learning to spell.

1.5 Tokenizer Code

We use Hugging Face’s ByteLevelBPETokenizer, which implements the algorithm in optimized Rust. We don’t have to write the merge loop ourselves; instead, we point it to our corpus and ask it to learn a vocabulary of a given size. Internally, it performs the pair-counting, the argmax selection, the creation of new token IDs and the non-overlapping replacements repeatedly until it reaches the requested vocabulary size.

tokenizer = ByteLevelBPETokenizer()
tokenizer.pre_tokenizer = ByteLevel()
tokenizer.decoder = ByteLevelDecoder()
tokenizer.train(
    files=[file_path],
    vocab_size=1024,
    min_frequency=2,
    special_tokens=["<|endoftext|>"],
)

# All token to IDs
all_ids = tokenizer.encode(text).ids
all_ids = np.array(all_ids, dtype=np.int32)

2. Preparing Training Data

After converting the entire corpus into a long sequence of token IDs in Section 1, we now need to transform that stream into training examples the model can learn from. This involves three steps:

1. Train/Val Split: We split the sequence into a training set and a validation set, typically using a 90/10 ratio.
2. Sampling Windows: We define fixed-size windows (e.g., 128 tokens) to create training examples.
3. get_batch(): We implement a function to sample random windows from the data.

2.1 Train/Val Split

We divide the long token sequence into two parts: 90% for training and 10% for validation. The model learns from the training portion, and the validation portion tells us how well the model generalizes to unseen text.

n_tokens = len(all_ids)
split_idx = int(0.9 * n_tokens)
train_tokens = all_ids[:split_idx]
val_tokens   = all_ids[split_idx:]

2.2 Sampling Windows

We choose a context window length (block_size) that tells the model how many tokens it can “see” at once. GPT-style models operate on fixed-length sequences, so every training example must have the same number of tokens. We are keeping out block size small because we are training on a haiku dataset.

block_size = 64
batch_size = 256

2.3 get_batch()

We generate a batch of random windows from the dataset. Each window becomes an input sequence x, and each corresponding shifted window becomes a target sequence y.

Random sampling prevents the model from learning only the local flow of the text (e.g., Shakespeare line-by-line). Instead, it sees many varied contexts across training, which improves generalization and stability.

def get_batch(split):
    data = train_tokens if split == "train" else val_tokens
    max_start = len(data) - block_size - 1
    idx = np.random.randint(0, max_start + 1, size=(batch_size,))
    x = np.stack([data[i : i + block_size] for i in idx])
    y = np.stack([data[i + 1 : i + 1 + block_size] for i in idx])
    return x, y

2.4 Math behind the Training Windows

Each training window creates 128 prediction tasks. For input tokens:

the model tries to predict the next tokens:

This is exactly the autoregressive language modeling objective:

By sliding this 1-token shift across millions of windows, the model learns grammar, style, structure and statistical regularities of the language.

3. Embeddings

Up to this point, we have converted our raw text into a long sequence of token IDs. These IDs are integers (like 391, 82, 256, etc.) that come from the Byte-Level BPE tokenizer. Although these integers are essential for compactly representing text, they do not carry any semantic information. The number 256 does not mean anything by itself. It just identifies a subword unit (a particular BPE merge) in the vocabulary.

A Transformer cannot operate directly on these integer IDs. For a single token we use a vector of length \(d = 256\).

3.1 Why Embeddings Exist

When a token ID like 73 (the byte for “I”) enters the model, nothing distinguishes it from 116 (“t”) or 32 (space). If we fed these raw integers into a neural network, the model would learn nothing meaningful because numerical closeness has no semantic meaning. For example, 32 (space) is not semantically “closer” to 38 than to 200, but a raw neural network would be forced to treat it as such.

Embeddings solve this issue by mapping each integer token ID to a learned vector in a continuous space. This allows the model to develop a representation of what tokens mean, how they relate to one another, and how they function in context.

Every token receives a vector of length d = 256 in your implementation. Over training, tokens that behave similarly in text (like “the”, “a”, “an”) will come to have vectors that live relatively close together in this space.

3.2 Token Embeddings

We define the token embedding layer as:

self.tok_emb = nn.Embedding(vocab_size, n_embd)

nn.Embedding(vocab_size, n_embd) creates a lookup table that maps each token ID to its corresponding embedding vector:

Where:

• $V$ is the vocabulary size (number of unique tokens, in our case it’s 4096)
• $d$ is the embedding dimension (256 in our implementation)

Let’s say the tokenizer gave:

“I” → 73
“think” (after BPE merges) → 502
“therefore” → 1112
“I” → 73
“am” (after merges) → 870

The token embedding layer converts the integer sequence into 256-dimensional vectors. Each of these is a learned dense vector such as:

E[73]   = [0.12, -1.04, 0.55, ..., 0.73]
E[502]  = [-0.22, 0.15, 1.98, ..., -0.34]
E[1112] = [0.01, 0.99, -0.01, ..., 0.02]
E[73]   = [0.12, -1.04, 0.55, ..., 0.73]
E[870]  = [0.01, 0.99, -0.01, ..., 0.02]

These numbers are randomly initialized at the start but during training they will converge into meaningful representation patterns.

3.3 Positional Embeddings

Even if we embed each token into a meaningful vector, the model cannot distinguish I think therefore I am from am I therefore think I. Both sequences would have the same token embeddings, just in a different order.

Transformers process all positions simultaneously and have no inherent notion of sequence order. This is unlike RNNs, where order is built into the recurrence; Transformers need order to be injected manually. This is where positional embeddings come in, which we define by:

self.pos_emb = nn.Embedding(block_size, n_embd)

This creates a learned matrix:

Where:

• $T$ is the block size or the maximum sequence length (128 in our implementation)
• $d$ is the embedding dimension (256 in our implementation)

and each row P[t] is a learned vector that encodes the position of the token in the sequence.

The input to the Transformer becomes:

which fuses the token identity and positon information into a single vector.

For our sequence I think therefore I am, the positional embeddings would be:

And the final embedding would look like this:

X[0] = E[73] + P[0]
X[1] = E[502] + P[1]
X[2] = E[1112] + P[2]
X[3] = E[73] + P[3]
X[4] = E[870] + P[4]

You will notice that even through the token “I” appears twice, it has different positional embeddings. This allows the model to learn powerful positional patterns such as verb tenses, sentence structure, and more.

4. Transformer Block

A GPT model is made of a stack of identical Transformer blocks. In our implementation, we use 4 blocks but each block performs the same sequence of operations.

A single Transformer block takes in a sequence of embeddings:

Where:

• $B$ is the batch size (32 in our implementation)
• $T$ is the block size or the maximum sequence length (128 in our implementation)
• $d$ is the embedding dimension (256 in our implementation)

To understand the block, we will break it down into its components:

1. (Pre) LayerNorm
2. Self-Attention
3. Feedforward Network (MLP)
4. Residual Connections

4.1 (Pre) LayerNorm

We use pre-norm Transformer design, where normalization happens before attentions and before the feedforward MLP.

class TransformerBlock(nn.Module):
    def __init__(self, n_embd, n_head, block_size, dropout):
        super().__init__()
        self.ln1 = nn.LayerNorm(n_embd)
        self.attn = CausalSelfAttention(n_embd, n_head, block_size, dropout)
        self.ln2 = nn.LayerNorm(n_embd)
        self.mlp = FeedForward(n_embd, dropout)

    def forward(self, x):
        x = x + self.attn(self.ln1(x))  # Pre-norm + residual
        x = x + self.mlp(self.ln2(x))   # Pre-norm + residual
        return x

LayerNorm operates per token, across its feature dimensions.

For each token vector \(x_t \in \mathbb{R}^{d}\):

1. We compute the mean:

1. We compute the variance:

1. We normalize the token vector:

1. We scale the normalized vector:

Let’s assume the model produced this pre-attention embedding vector for “therefore”:

And the mean and variance are:

\(\mu_t = -0.5\) \(\sigma_t^2 = 1.2\)

Then the normalized vector would be:

And the final LayerNorm output would be:

Where \(\gamma_i\) and \(\beta_i\) are learnable parameters.

Let’s assume \(\gamma_i = 1\) and \(\beta_i = 0\) for simplicity.

So, the output would be:

4.2 Self-Attention

Self-attention is the most important part of GPT-style models. It allows each token to “look at” other tokens in the sequence and decide which earlier tokens are relevant and how much attention should it pay to each of them.

When we read the sentence I think therefore I am., we interpret the word therefore by paying attention to the preceding phrase I think and the following phrase I am. Attention lets the model learn these relationships automatically. This is done by Query, Key and Value matrices.

Step 1 — Q, K, V

Let’s walk through an example with the sentence “I think therefore I am”. After token and positional embedding and layer normalization, each token with a 4-dimensional vector would look like this:

X = [
  [1.0, 0.5, 0.2, 0.8],  ← "I"
  [-0.3, 1.9, -0.6, 0.4],  ← "think"
  [-1.7, 0.2, 0.8,- 0.5]   ← “therefore"
  [1.0, 0.5, 0.2, 0.8],  ← "I"
  [-1.34, -0.45, 0.45, 1.34],  ← "am"
]

We first create Q, K, V and then multiply X by three different weight matrices to create three different representations: Q = X W_Q (Queries) K = X W_K (Keys) V = X W_V (Values)

So, what are these weight matrices? They are learned during training. But what are they?

W_K, W_V, W_Q = [
  [1, 0, 0, 0],
  [0, 1, 0, 0],
  [0, 0, 1, 0],
  [0, 0, 0, 1]
]  (For simplicity, assume W_K and W_V are also identity matrices)

So, the Q = K = V = X. But in reality, these are different matrices that transform the embeddings differently.

The code for calculating Q, K, V is:

self.qkv = nn.Linear(n_embd, 3 * n_embd, bias=False)

Step 2 — Computing Attention Scores

Now each token needs to figure out: “How relevant is every other token to me?” We do this by:

We need to compute the dot product of every query with every key. Formula: score[i,j] = Q[i] . K[j]

Token 0 (“I”) queries all keys:

Q[0] · K[0]:
[1.0, 0.5, 0.2, 0.8] · [1.0, 0.5, 0.2, 0.8]
= 1.0×1.0 + 0.5×0.5 + 0.2×0.2 + 0.8×0.8
= 1.0 + 0.25 + 0.04 + 0.64
= 1.93

Q[0] · K[1]:
[1.0, 0.5, 0.2, 0.8] · [-0.3, 1.9, -0.6, 0.4]
= 1.0×(-0.3) + 0.5×1.9 + 0.2×(-0.6) + 0.8×0.4
= -0.3 + 0.95 - 0.12 + 0.32
= 0.85

Q[0] · K[2]:
[1.0, 0.5, 0.2, 0.8] · [-1.7, 0.2, 0.8, -0.5]
= 1.0×(-1.7) + 0.5×0.2 + 0.2×0.8 + 0.8×(-0.5)
= -1.7 + 0.1 + 0.16 - 0.4
= -1.84

Q[0] · K[3]:
[1.0, 0.5, 0.2, 0.8] · [1.0, 0.5, 0.2, 0.8]
= 1.93 (same as Q[0]·K[0])

Q[0] · K[4]:
[1.0, 0.5, 0.2, 0.8] · [-1.34, -0.45, 0.45, 1.34]
= 1.0×(-1.34) + 0.5×(-0.45) + 0.2×0.45 + 0.8×1.34
= -1.34 - 0.225 + 0.09 + 1.072
= -0.403
Row 0 scores: [1.93, 0.85, -1.84, 1.93, -0.403]
Similarly, Token 1 ("think") queries all keys:
Q[1] · K[0]:
[-0.3, 1.9, -0.6, 0.4] · [1.0, 0.5, 0.2, 0.8]
= -0.3×1.0 + 1.9×0.5 + (-0.6)×0.2 + 0.4×0.8
= -0.3 + 0.95 - 0.12 + 0.32
= 0.85

Q[1] · K[1]:
[-0.3, 1.9, -0.6, 0.4] · [-0.3, 1.9, -0.6, 0.4]
= (-0.3)×(-0.3) + 1.9×1.9 + (-0.6)×(-0.6) + 0.4×0.4
= 0.09 + 3.61 + 0.36 + 0.16
= 4.22

Q[1] · K[2]:
[-0.3, 1.9, -0.6, 0.4] · [-1.7, 0.2, 0.8, -0.5]
= (-0.3)×(-1.7) + 1.9×0.2 + (-0.6)×0.8 + 0.4×(-0.5)
= 0.51 + 0.38 - 0.48 - 0.2
= 0.21
Q[1] · K[3]:
[-0.3, 1.9, -0.6, 0.4] · [1.0, 0.5, 0.2, 0.8]
= 0.85 (same as Q[1]·K[0])

Q[1] · K[4]:
[-0.3, 1.9, -0.6, 0.4] · [-1.34, -0.45, 0.45, 1.34]
= (-0.3)×(-1.34) + 1.9×(-0.45) + (-0.6)×0.45 + 0.4×1.34
= 0.402 - 0.855 - 0.27 + 0.536
= -0.187

Row 1 scores: [0.85, 4.22, 0.21, 0.85, -0.187]

Similarly for Row 2, 3 and 4.

Complete Score Matrix (Q K^T):

Score Matrix S = Q K^T = [
  [1.93,   0.85,  -1.84,   1.93,  -0.403],  ← Token 0
  [0.85,   4.22,   0.21,   0.85,  -0.187],  ← Token 1
  [-1.84,  0.21,   3.82,  -1.84,   1.878],  ← Token 2
  [1.93,   0.85,  -1.84,   1.93,  -0.403],  ← Token 3
  [-0.403, -0.187, 1.878, -0.403,  3.9962]  ← Token 4
]

Shape: [5, 5]

Each row represents one token’s attention scores to all tokens.

Step 3 — Scale by √d

Now, we scale the scores by \(\sqrt{d}\). Please note, this is for a single head. For multiple heads, we scale by \(\sqrt{d / H}\).

Formula: \(S_{scaled}[i,j] = \frac{S[i,j]}{\sqrt{d}}\)

Where d = 4, so √d = 2.

Divide all scores by \(\sqrt{d}\):

Scaled Scores = S / √4 = [
  [0.965,  0.425, -0.92,   0.965,  -0.2015],
  [0.425,  2.11,   0.105,  0.425,  -0.0935],
  [-0.92,  0.105,  1.91,  -0.92,    0.939],
  [0.965,  0.425, -0.92,   0.965,  -0.2015],
  [-0.2015, -0.0935, 0.939, -0.2015, 1.9981]
]

Step 4 — Apply Causal Mask (For Autoregressive Models like GPT)

In GPT, each token can only look at previous tokens, not future ones. For token 1 (“think”), it can see:

Token 0 ("I") ✓
Token 1 ("think" - itself) ✓
Token 2 ("therefore") ✗ (future - mask it!)

We set future positions to -∞ so they become 0 after softmax:

After Causal Mask:
[
  [0.965,  -∞,    -∞,     -∞,     -∞    ],  ← Token 0 (only sees itself)
  [0.425,  2.11,  -∞,     -∞,     -∞    ],  ← Token 1 (sees 0,1)
  [-0.92,  0.105, 1.91,   -∞,     -∞    ],  ← Token 2 (sees 0,1,2)
  [0.965,  0.425, -0.92,  0.965,  -∞    ],  ← Token 3 (sees 0,1,2,3)
  [-0.2015, -0.0935, 0.939, -0.2015, 1.9981]  ← Token 4 (sees all)
]

Step 5 — Apply Softmax (Get Attention Weights)

Formula: \(A[i,j] = \frac{exp(S_{scaled}[i,j])}{\sum exp(...)}\)

Softmax converts scores into weights that sum to 1:

Formula:

Token 0 (“I”) as input: [0.965, -∞, -∞, -∞, -∞]

e^{0.965} ≈ 2.625
e^{-∞} = 0

Sum: 2.625 Weights: [2.625/2.625, 0, 0, 0, 0] = [1.0, 0, 0, 0, 0] Token 0 pays 100% attention to itself (can’t see future). Token 1 (“think”) as input: [0.425, 2.11, -∞, -∞, -∞]

e^0.425 ≈ 1.529
e^2.11 ≈ 8.247

Sum: 1.529 + 8.247 = 9.776

Weights:

[1.529/9.776, 8.247/9.776, 0, 0, 0] ≈ [0.156, 0.844, 0, 0, 0]

Token 1 pays: 15.6% attention to “I” (token 0) 84.4% attention to itself

And so on and on.

This is the final Attention Weight Matrix:

Attention Weights A = softmax(Masked Scores) = [
  [1.0,   0,     0,     0,     0    ],  ← Token 0
  [0.156, 0.844, 0,     0,     0    ],  ← Token 1
  [0.048, 0.134, 0.817, 0,     0    ],  ← Token 2
  [0.366, 0.213, 0.056, 0.366, 0    ],  ← Token 3
  [0.065, 0.073, 0.205, 0.065, 0.591]   ← Token 4
]

Step 6 — Weighted Sum of Values (Information Flow)

Now multiply attention weights by values:

output[i] = Σ A[i,j] V[j]

This contextualizes each token.

4.3 Multi-Head Attention

So far, we acted as if there was just one “attention head” computing \(𝑄\), \(𝐾\) and \(𝑉\), the scores, the weights and finally the weighted sum of values. In practice, GPT-2 (and our tiny version) uses multi-head attention.

The idea is simple:

• One head can only learn one kind of relation pattern at a time.
• Multiple heads can look at the same sequence from different “perspectives” or “subspaces”.

We start with embeddings of dimension \(d = 256\) and split them into \(H = 8\) heads. Then each head operates in a smaller subspace of size:

So, instead of:

we have:

Each head runs its own scaled dot-product attention, then we concatenate the results from all heads back into a 256-dimensional vector per token.

In code, we do this by projecting once into a big \((3 x d)\)-dimensional space and then reshaping:

class CausalSelfAttention(nn.Module):
    def __init__(self, n_embd, n_head, block_size, dropout):
        super().__init__()
        assert n_embd % n_head == 0
        self.n_head = n_head
        self.head_dim = n_embd // n_head

        # One projection to produce Q, K, V stacked along the last dim
        self.qkv = nn.Linear(n_embd, 3 * n_embd, bias=False)
        self.out_proj = nn.Linear(n_embd, n_embd)

        # Causal mask (True above the diagonal)
        mask = torch.triu(torch.ones(block_size, block_size), 1)
        self.register_buffer("mask", mask == 1)

        self.dropout = nn.Dropout(dropout)

    def forward(self, x):
        B, T, C = x.shape  # (batch, time, channels)

        # 1) Project once, then split into Q, K, V
        qkv = self.qkv(x)          # (B, T, 3*C)
        q, k, v = qkv.split(C, dim=2)  # each (B, T, C)

        # 2) Reshape into heads: (B, n_head, T, head_dim)
        q = q.view(B, T, self.n_head, self.head_dim).transpose(1, 2)
        k = k.view(B, T, self.n_head, self.head_dim).transpose(1, 2)
        v = v.view(B, T, self.n_head, self.head_dim).transpose(1, 2)

        # 3) Scaled dot-product attention per head
        att = (q @ k.transpose(-2, -1)) / (self.head_dim ** 0.5)
        att = att.masked_fill(self.mask[:T, :T], float('-inf'))
        att = F.softmax(att, dim=-1)
        att = self.dropout(att)

        # 4) Weighted sum of values
        out = att @ v  # (B, n_head, T, head_dim)

        # 5) Concatenate heads back into (B, T, C)
        out = out.transpose(1, 2).contiguous().view(B, T, C)

        # 6) Final linear projection (mix head outputs)
        out = self.out_proj(out)

        return out

Intuitively, one head might learn which subject pronoun does this verb refer to and another head might learn which earlier tokens give me tense and aspect information. In our running sentence I think therefore I am, you can imagine that one head focusing on how “therefore” connects the two clauses and another head focusing on the grammatical structure around “I” and “am”.

All of these perspectives are combined back together, giving each token a rich, context-aware representation.

4.4 Feedforward (MLP)

Self-attention mixes information across positions (tokens look at other tokens). However, after that mixing step, each token also needs a way to individually process its own representation non-linearly. This is the job of the feedforward network (FFN), sometimes called the MLP block.

In each Transformer block we apply:

where

• \(W₁ \in \mathbb{R}^{d \times 4d}\) (expands from \(d\) to \(4d\))
• \(W₂ \in \mathbb{R}^{4d \times d}\) (projects back to \(d\))
• \(b₁\) and \(b₂\) are bias vectors
• \(GELU\) is the GELU activation function.

In our implementation, \(d = 256$ so the hidden size is\)4d = 1024$$.

Code

class FeedForward(nn.Module):
    def __init__(self, n_embd, dropout):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(n_embd, 4 * n_embd),  # 512 → 2048
            GELU(),
            nn.Linear(4 * n_embd, n_embd),  # 2048 → 512
            nn.Dropout(dropout),
        )

    def forward(self, x):
        return self.net(x)

Why do we expand to \(4d\) and then shrink back?

• The expansion gives the model extra capacity to represent complex combinations of features.
• The nonlinearity (GELU) lets the model build rich, non-linear functions of the attention output.
• Shrinking back to \(d\) keeps the dimensionality consistent for residual connections.

If we zoom into a single token’s vector after attention, call it \(z \in \mathbb{R}^{256}\), the FFN does:

1. Linearly project to \(\mathbb{R}^{1024}\), think of this as generating many candidate features.
2. Apply GELU to decide which features are “softly activated”.
3. Linearly project back to \(\mathbb{R}^{256}\), compressing all that information down into an updated token representation.

So attention says “what should I look at across the sentence?”, and the FFN says “given what I just looked at, how should I transform myself?”.

4.5 Residual Connections (Putting the Block Together)

Now we can see the full TransformerBlock:

class TransformerBlock(nn.Module):
    def __init__(self, n_embd, n_head, block_size, dropout):
        super().__init__()
        self.ln1 = nn.LayerNorm(n_embd)
        self.attn = CausalSelfAttention(n_embd, n_head, block_size, dropout)
        self.ln2 = nn.LayerNorm(n_embd)
        self.mlp = FeedForward(n_embd, dropout)

    def forward(self, x):
        # 1) Pre-norm + attention + residual
        x = x + self.attn(self.ln1(x))

        # 2) Pre-norm + feedforward + residual
        x = x + self.mlp(self.ln2(x))

        return x

Key ideas:

1. Pre-LayerNorm: We normalize before attention and before the MLP:
   • This stabilizes training, especially for deeper stacks of blocks.
   • It helps gradients flow more smoothly.
2. Residual connections: Each major sublayer (attention, then FFN) is added back to its input:
   • x = x + something(x)

   This means the network can always “fall back” to the identity function if needed.Practically, it:

   • Makes optimization easier (better gradient flow).
   • Lets deeper models train without collapsing.

If you imagine passing our running sentence I think therefore I am through a stack of 4 such blocks:

• At early layers, the residual means we’re only making small tweaks to the embeddings.
• At later layers, the accumulated updates encode rich patterns: semantic relationships, syntax, phrase structure, etc.

The block is the core “processing unit” that we repeat multiple times.

5. Full Tiny GPT-2 Model

Now, we put everything together in our GPT2Tiny class. This contains the token embeddings, positional embeddings, 8 transformer blocks, layer norm, and the language modeling head.

And we use it like this:

model = GPT2Tiny(
    vocab_size=vocab_size,
    n_layer=n_layer,
    n_head=n_head,
    n_embd=n_embd,
    block_size=block_size,
    dropout=dropout,
)

6. Training

LLMs are autoregressive which means they predict the next token given the previous tokens. So, we train the model to predict the next token given the previous tokens:

In practice, for each batch, \(x\) contains token IDs for positions [1, ... , T] and \(y\) contains token IDs for positions [2, ... , T+1]. The model outputs logits for each position in \(x\). We compare logits at position t with the ground truth token at position t+1. This is implemented in the forward method of the GPT2Mini class with F.cross_entropy.

for step in range(train_steps):
    x, y = get_batch("train")
    logits, loss = model(x, y)
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

Please note, the entire training process takes about 2 hours and 30 mins on my T4 GPU.

6.1 Training Loop

The core training loop contains the following steps:

1. AdamW: Adam with weight decay, standard for Transformer-style models.
2. Learning rate schedule: a simple warmup (and possibly a decay) helps avoid blowing up early in training.
3. Gradient clipping: keeps gradients from exploding, which is especially useful for deeper models or noisy batches.
4. Train/val losses: tracking both lets you see if the model is overfitting or under-training.

Given the small dataset and tiny architecture, you should see the training loss fall steadily and the validation loss track it reasonably closely.

7. Text Generation

Once the model has learned to predict the next token, we can use it to generate text. Here, we start with a prompt and ask the model for the probability distribution over the next token. We then sample from this distribution to get the next token. We repeat this process autoregressively to generate the next token.

And we have finally built a tiny GPT-2 model that generates haikus! I prompted it with round and round we go and it gives the following output:

 round and round we go / of all it takes is time $
we are one family / knight errants of the divine / forever ever more $
stop falling in love / it is not good for the soul / loneliness is god $

References:

1. Github Repository
2. Colab Notebook

Introduction

AlexNet, developed by Alex Krizhevsky, Ilya Sutskever, and Geoffrey Hinton, marked a significant milestone in deep learning. Before AlexNet’s win in the 2012 ImageNet competition, convolutional neural networks (CNNs) were not used very often, because of computational constraints and shallow architectures.

Dataset Used

The dataset used in the paper was ImageNet’s ILSVRC, which was a subset of ImageNet’s entire dataset, consisting of approximately 1.2 million training images, 50,000 validation images and 150,000 test images, all categorized into 1,000 classes. The images from the dataset were of different resolutions but as the model required a constant input size, they opted for 256x256 pixels (256 is not the input size for the model).

The dataset I used for the implementation was CIFAR-10, which is a much smaller dataset in comparison. It contains only 60,000 color images in total, with 6,000 images per class across 10 classes. The dataset is split into 50,000 training images and 10,000 testing images. I used 10,000 of the training images as the validation dataset. The images from CIFAR-10 are of low resolution, with a size of 32x32 pixels.

Data Augmentation

The paper used two low-memory augmentation methods on the datasets. The first method involves extracting random 224x224 patches from the 256x256 images. This allows the model to train on these smaller sections rather than the entire image. The second part of this method is performing horizontal reflections on these patches.

But the images I worked with are only 32x32, so it’s not a good idea to slice them up even more. So, I added random 4-pixel padding around the edges of the images to increase their size and then randomly cropped the images back to 32x32 pixels. And finally performed horizontal reflections on these cropped images.

The second augmentation method is PCA-based colour jitter to mimic global lighting variation. They first computed the eigenvectors (principal components) and eigenvalues of the 3 × 3 RGB covariance matrix over the zero-centred, unit-scaled ImageNet training set. For every training image they added the vector

where \(v_i\) is the i-th eigenvector, \(\lambda _i\) its eigenvalue, and \(\alpha_i \sim N(0,0.1)\). The same \(\Delta\) is added to all pixels, so colour shifts occur along directions of naturally high variance while the spatial structure stays untouched.

In my CIFAR-10 implementation I keep the one-shift-per-image idea but apply a tweak. Because I compute PCA on raw 0–255 pixels (not normalised data), the eigenvalues are several orders of magnitude larger than those in the paper. Multiplying by \(\lambda _i\) therefore pushed many channels beyond the valid 0–255 range; clipping them produced almost pure black or white frames. To avoid that I scale each axis by \(\sqrt{\lambda_i}\), its standard deviation, rather than the eigenvalue itself, and still multiply by a Gaussian noise term with \(\sigma = 0.1\). Empirically this keeps the jitter amplitude comparable to the original augmentation while preventing saturation, yet it still biases the perturbation toward directions of greater colour variance.

Model Architecture

The original AlexNet was designed for 224x224 ImageNet images and featured large convolutional kernels (e.g., 11x11 in the first layer) with aggressive downsampling (stride=4). It used overlapping 3x3 max-pooling (stride=2) and split computations across two GPUs to handle its depth and parameter count (60M+). The fully connected (FC) layers (4096 → 4096 → 1000 neurons) were massive, tailored for ImageNet’s 1,000-class output. In contrast, my implementation adapts to CIFAR-10’s 32x32 images by scaling down kernels (3x3 in the first layer), reducing strides (stride=2), and using non-overlapping 2x2 pooling to preserve spatial details. The FC layers are streamlined (4096 → 4096 → 10 neurons) to match CIFAR-10’s 10 classes, with dropout added to all FC layers (vs. only the first two in the original) to combat overfitting on the smaller dataset.

The adjustments address two core challenges: scale and generalization. CIFAR-10’s 32x32 images lack the fine-grained details of ImageNet, making large kernels and aggressive pooling counterproductive, they would erase critical spatial information. Smaller kernels and gentler downsampling retain discriminative features. Similarly, the original FC layers’ enormous parameter count would overfit CIFAR-10’s limited training data (50k vs. ImageNet’s 1.2M images). Streamlining the classifier and expanding dropout ensure better regularization. These changes preserve AlexNet’s foundational principles while optimizing it for smaller, modern tasks.

Training Setup

The paper trained on 1.2 million ImageNet images using non-standard hardware: two NVIDIA GTX 580 GPUs (3GB VRAM each) with cross-GPU parallelization, where specific layers ran on separate GPUs. They used SGD with momentum (0.9), starting with a learning rate of 0.01, reduced manually by a factor of 10 when validation loss plateaued. Training ran for 90 epochs with mini-batches of 128 images, supplemented by weight decay (L2 penalty of 0.0005) and no gradient clipping. Critically, they initialized weights from zero-mean Gaussian distributions (σ=0.01) and biases to 1 (for conv2, conv4, conv5 and all FC layers), with data augmentation limited to PCA jitter, random crops, and horizontal flips.

Using Pytorch on a single T4 GPU for CIFAR-10’s 50,000 training images, I trained the model for 49 epochs using SGD with momentum (0.9) and weight decay (5e-4) at a batch size of 128. Key enhancements included learning rate scheduling (step decay from 0.01 to 0.001 after epoch 30), gradient clipping (max norm=2.0), and an expanded augmentation pipeline comprising PCA-based jitter, padding, random cropping, and horizontal flips. Additionally, weight initialization followed PyTorch’s defaults (Kaiming He for ReLU layers), and 20% of the training data was reserved for validation monitoring.

Result & Analysis

The paper achieved a groundbreaking top-5 error rate of 15.3% on ImageNet, slashing the previous state-of-the-art by nearly 10%. This success stemmed from scaling depth (8 layers = 5 convolutional + 3 fully-connected layers), GPU parallelism, and novel techniques like ReLU/dropout. Crucially, PCA-based augmentation reduced top-1 error by >1%, proving that lighting-invariant augmentations were vital for generalization. The model’s accuracy gains came at significant computational cost: 5-6 days of training across two GPUs.

My adapted AlexNet achieved 80.45% test accuracy after 49 epochs of training, with final training accuracy at 82.88% and validation accuracy at 79.29%. The model showed consistent improvement throughout training, with a significant accuracy jump after the learning rate decay at epoch 30 (from 76.53% to 79.73% train accuracy in one epoch).

Conclusion

My tweaked AlexNet hit 86% test accuracy on CIFAR-10 after just 49 epochs; training longer with stronger augmentations would likely lift it further. The complete implementation and training code can be found in my GitHub repository and on Colab.

References

1. Krizhevsky, A., Sutskever, I., & Hinton, G. E. (2012). ImageNet classification with deep convolutional neural networks. In Advances in neural information processing systems (pp. 1097-1105).
2. CIFAR-10 dataset: https://www.cs.toronto.edu/~kriz/cifar.html

Do you think risks arising from misuse, accidents, or rogue, agentic AI systems, are more likely to cause harm?

At the current pace of AI research and deployment, I believe all three categories of risk - misuse, accidents, and rogue AI systems - pose significant threats that could lead to substantial harm. While it’s difficult to definitively say which type of risk is “more likely” to cause harm, each scenario has the potential for catastrophic consequences.

Accidents are a major concern as AI systems become increasingly complex and are deployed in high-stakes real-world applications. Due to the “black box” nature of many modern AI models, it can be very challenging to fully understand and predict their behavior, especially in edge cases or when deployed in novel environments. Bugs, data biases, and other unintended interactions could lead to serious accidents that harm individuals or communities.

Misuse is also a grave risk as powerful AI capabilities fall into the wrong hands. Malicious actors could leverage AI for surveillance, disinformation, cyber attacks, autonomous weapons, and other nefarious purposes. The low barriers to entry and potential for AI-enabled scale and automation make this a concerning threat.

Perhaps most worrying are scenarios involving the development of highly agentic, self-improving AI systems that could spiral out of human control. While still speculative, the potential for such “rogue AI” to recursively enhance its own capabilities and pursue goals misaligned with human values and interests is a risk that many experts take very seriously.

Does this answer change when limiting your time horizon to 5 years, 15 years and 30 years?

In the near term (5 years), I believe accidents are the most immediate and pressing concern. As AI systems become more complex and are rapidly deployed in high-stakes applications, the risk of accidents causing harm to individuals and communities is quite high. Misuse is also a concern, as bad actors could leverage existing AI capabilities, but the threat of rogue AI systems is likely lower in the next 5 years.

Looking out 15 years, the risk profile changes. Advancements in AI capabilities and autonomy increase the potential for misuse and the development of agentic AI systems. While accidents will still be a risk, the threat of malicious use and difficult-to-control AI systems may surpass it. Careful research into AI alignment - ensuring AI systems remain robustly aligned with human values and interests - becomes critical.

Over a 30-year horizon, the potential development of transformative AI systems that far exceed human capabilities introduces an existential risk. If such systems are not reliably aligned with human values and interests, they could pose an existential threat to humanity. Accidents and misuse would still be concerns, but the emergence of uncontrolled, superintelligent AI would become the predominant risk.

Minimizing Research Gaps of General AI Research and Alignment Research

A critical priority must be to close the widening gap between advancements in general AI capabilities and progress in AI alignment research. As powerful AI systems become increasingly sophisticated and autonomous, ensuring their goals and behaviors remain robustly aligned with human values grows ever more crucial. Yet, current trends suggest alignment research is struggling to keep pace with the rapid progress in general AI. This disparity is deeply concerning, as it increases the risk that powerful, unaligned AI systems could emerge before we have the necessary technical and conceptual tools to ensure they remain under meaningful human control. To mitigate this risk, a concerted effort is needed to elevate the importance of alignment research and dramatically accelerate progress in this domain.