HeyMSK
Chapter 4 of 10

Build the network

We have token IDs and targets. Now I need a function that takes those IDs and scores the possible next tokens. I start by looking up a list of 128 adjustable numbers for each token. That list is called a vector. I add another vector to mark its position.

Attention mixes information from this position and earlier positions. A small neural network called a feed-forward layer transforms the result. I group these operations into a Transformer block and apply two blocks in sequence.

The final layer turns the resulting vector into 98 next-token scores. These calculations work before training, but the random weights have not learned which answer is right.

Follow the equals sign through the four operations. In operation 2, click the later 2 to see why it is hidden. In operation 4, compare the chosen token with the training target. Then score your own prefix below.

Follow one position through the network

Start at = in 1+1=2. Its training target is the following 2. Work through the four operations, then choose another position.

Turn the token ID into a vector

= has ID 32. That ID selects a row of 128 adjustable numbers. A second table supplies numbers for position 4. I add the two rows.

Vector1234
Token0.0180.0530.032-0.024
Position0.0120.010-0.013-0.011
Sum0.0290.0620.019-0.035

First 4 of 128 numbers, rounded to three decimals. Compare the two 1 positions: their token rows match, but their position rows differ. The network can tell them apart.

Recorded values for the fixed example, computed with the same untrained weights used in the live exercise below.

Build this step

The network carries an array shaped [batch, positions, 128]. A vector is a list of 128 numbers describing one token. An embedding is the learned table that supplies that vector.

Each layer below is a Python class based on nn.Module. Its __init__ method creates the weights. Its forward method defines the calculation. self refers to that layer.

self.token_embedding = nn.Embedding(98, 128)
self.position_embedding = nn.Embedding(32, 128)
# In the forward pass:
positions = torch.arange(input_ids.shape[1], device=input_ids.device)
x = self.token_embedding(input_ids) + self.position_embedding(positions)

Let earlier characters affect a prediction

Attention creates three versions of each vector. A query asks what is useful, a key describes what a position offers, and a value holds the information to combine. Matrix multiplication, written @, scores their matches.

class CausalSelfAttention(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.n_heads = config.n_heads
        self.head_width = config.width // config.n_heads
        self.qkv = nn.Linear(config.width, 3 * config.width)
        self.projection = nn.Linear(config.width, config.width)
        self.attention_dropout = nn.Dropout(config.dropout)
        self.output_dropout = nn.Dropout(config.dropout)
        future_positions = torch.ones(config.context_length, config.context_length,
                                     dtype=torch.bool).triu(diagonal=1)
        self.register_buffer("future_positions", future_positions, persistent=False)

    def forward(self, x):
        batch, length, width = x.shape
        q, k, v = self.qkv(x).chunk(3, dim=-1)
        q, k, v = [part.reshape(batch, length, self.n_heads, self.head_width).transpose(1, 2)
                   for part in (q, k, v)]
        scores = (q @ k.transpose(-2, -1)) / math.sqrt(self.head_width)
        scores = scores.masked_fill(self.future_positions[:length, :length], float("-inf"))
        probabilities = self.attention_dropout(scores.softmax(dim=-1))
        mixed = (probabilities @ v).transpose(1, 2).contiguous().reshape(batch, length, width)
        return self.output_dropout(self.projection(mixed))

The four heads do this in parallel with 32 numbers each. Reshape separates the heads, and transpose arranges the array so matrix multiplication compares token positions. softmax turns scores into weights that sum to one. The triangular mask gives future positions zero weight. Otherwise the network could read the answer during training.

Transform each position and keep its earlier information

class FeedForward(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.expand = nn.Linear(config.width, config.feed_forward_width)
        self.activation = nn.GELU()
        self.contract = nn.Linear(config.feed_forward_width, config.width)
        self.dropout = nn.Dropout(config.dropout)

    def forward(self, x):
        return self.dropout(self.contract(self.activation(self.expand(x))))

class TransformerBlock(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.attention_norm = nn.LayerNorm(config.width)
        self.attention = CausalSelfAttention(config)
        self.feed_forward_norm = nn.LayerNorm(config.width)
        self.feed_forward = FeedForward(config)

    def forward(self, x):
        x = x + self.attention(self.attention_norm(x))
        return x + self.feed_forward(self.feed_forward_norm(x))

The feed-forward layer expands 128 numbers to 512, applies a nonlinear function called GELU, then returns to 128. A residual addition keeps the old vector and adds the new result. Layer normalization keeps the numerical scale controlled.

I repeat the block twice and turn the final vectors into 98 token scores. The output reuses the token embedding weights. The complete file below includes the configuration, initialization and final output layer.

Read the complete model.py
"""A small GPT-style causal Transformer, initialized without pretrained weights."""

from dataclasses import dataclass
import math

import torch
from torch import nn
from torch.nn import functional as F

from .tokenizer import CharacterTokenizer


@dataclass(frozen=True)
class ModelConfig:
    vocab_size: int = 98
    context_length: int = 32
    n_layers: int = 2
    width: int = 128
    n_heads: int = 4
    feed_forward_width: int = 512
    dropout: float = 0.0

    def __post_init__(self):
        for name in ("vocab_size", "context_length", "n_layers", "width", "n_heads", "feed_forward_width"):
            value = getattr(self, name)
            if type(value) is not int or value <= 0:
                raise ValueError(f"{name} must be a positive integer")
        if self.vocab_size != CharacterTokenizer().vocab_size:
            raise ValueError("Model vocabulary must match printable-ascii-v1's 98 tokens")
        if self.width % self.n_heads:
            raise ValueError("Model width must divide evenly among attention heads")
        if type(self.dropout) not in (int, float) or not 0 <= self.dropout < 1:
            raise ValueError("Dropout must be between 0 (inclusive) and 1 (exclusive)")


class CausalSelfAttention(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.n_heads = config.n_heads
        self.head_width = config.width // config.n_heads
        self.qkv = nn.Linear(config.width, 3 * config.width)
        self.projection = nn.Linear(config.width, config.width)
        self.attention_dropout = nn.Dropout(config.dropout)
        self.output_dropout = nn.Dropout(config.dropout)
        future_positions = torch.ones(config.context_length, config.context_length,
                                     dtype=torch.bool).triu(diagonal=1)
        self.register_buffer("future_positions", future_positions, persistent=False)

    def forward(self, x):
        batch, length, width = x.shape
        q, k, v = self.qkv(x).chunk(3, dim=-1)
        q, k, v = [part.reshape(batch, length, self.n_heads, self.head_width).transpose(1, 2)
                   for part in (q, k, v)]
        scores = (q @ k.transpose(-2, -1)) / math.sqrt(self.head_width)
        scores = scores.masked_fill(self.future_positions[:length, :length], float("-inf"))
        probabilities = self.attention_dropout(scores.softmax(dim=-1))
        mixed = (probabilities @ v).transpose(1, 2).contiguous().reshape(batch, length, width)
        return self.output_dropout(self.projection(mixed))


class FeedForward(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.expand = nn.Linear(config.width, config.feed_forward_width)
        self.activation = nn.GELU()
        self.contract = nn.Linear(config.feed_forward_width, config.width)
        self.dropout = nn.Dropout(config.dropout)

    def forward(self, x):
        return self.dropout(self.contract(self.activation(self.expand(x))))


class TransformerBlock(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.attention_norm = nn.LayerNorm(config.width)
        self.attention = CausalSelfAttention(config)
        self.feed_forward_norm = nn.LayerNorm(config.width)
        self.feed_forward = FeedForward(config)

    def forward(self, x):
        x = x + self.attention(self.attention_norm(x))
        return x + self.feed_forward(self.feed_forward_norm(x))


class TinyArithmeticGPT(nn.Module):
    def __init__(self, config=None):
        super().__init__()
        self.config = config or ModelConfig()
        self.token_embedding = nn.Embedding(self.config.vocab_size, self.config.width)
        self.position_embedding = nn.Embedding(self.config.context_length, self.config.width)
        self.embedding_dropout = nn.Dropout(self.config.dropout)
        self.blocks = nn.ModuleList(TransformerBlock(self.config) for _ in range(self.config.n_layers))
        self.final_norm = nn.LayerNorm(self.config.width)
        self.lm_head = nn.Linear(self.config.width, self.config.vocab_size, bias=False)
        self.apply(self._initialize)
        # One learned table is used for both input representations and output scores.
        self.lm_head.weight = self.token_embedding.weight
        # Scale projections on residual branches as in the GPT-2 initialization recipe.
        residual_std = 0.02 / math.sqrt(2 * self.config.n_layers)
        for block in self.blocks:
            nn.init.normal_(block.attention.projection.weight, mean=0.0, std=residual_std)
            nn.init.normal_(block.feed_forward.contract.weight, mean=0.0, std=residual_std)

    @staticmethod
    def _initialize(module):
        if isinstance(module, (nn.Linear, nn.Embedding)):
            nn.init.normal_(module.weight, mean=0.0, std=0.02)
            if isinstance(module, nn.Linear) and module.bias is not None:
                nn.init.zeros_(module.bias)

    def hidden_states(self, input_ids):
        if not isinstance(input_ids, torch.Tensor) or input_ids.dtype != torch.long:
            raise ValueError("Input IDs must be a torch.long tensor")
        if input_ids.ndim != 2 or min(input_ids.shape) == 0:
            raise ValueError("Input IDs must have nonempty batch and sequence dimensions")
        if input_ids.shape[1] > self.config.context_length:
            raise ValueError("Input sequence exceeds the context length; no tokens were truncated")
        if bool(((input_ids < 0) | (input_ids >= self.config.vocab_size)).any()):
            raise ValueError("Input contains an invalid token ID")
        positions = torch.arange(input_ids.shape[1], device=input_ids.device)
        x = self.token_embedding(input_ids) + self.position_embedding(positions)
        x = self.embedding_dropout(x)
        for block in self.blocks:
            x = block(x)
        return self.final_norm(x)

    def forward(self, input_ids):
        # Logits are raw scores, not probabilities. Keep them for stable loss calculation.
        return self.lm_head(self.hidden_states(input_ids))

    def parameter_breakdown(self):
        counts = {
            "token_embedding_and_shared_output": self.token_embedding.weight.numel(),
            "position_embedding": self.position_embedding.weight.numel(),
            "transformer_blocks": sum(parameter.numel() for parameter in self.blocks.parameters()),
            "final_normalization": sum(parameter.numel() for parameter in self.final_norm.parameters()),
        }
        counts["total"] = sum(parameter.numel() for parameter in self.parameters())
        return counts


def next_token_loss(logits, target_ids):
    if (not isinstance(target_ids, torch.Tensor) or target_ids.dtype != torch.long
            or logits.ndim != 3 or logits.shape[-1] != CharacterTokenizer().vocab_size
            or logits.shape[:2] != target_ids.shape):
        raise ValueError("Targets must be long IDs matching logits' batch and sequence dimensions")
    if logits.device != target_ids.device:
        raise ValueError("Logits and targets must be on the same device")
    if bool(((target_ids < 0) | (target_ids >= logits.shape[-1])).any()):
        raise ValueError("Target contains an invalid token ID")
    if not bool((target_ids != CharacterTokenizer.PAD_ID).any()):
        raise ValueError("A loss needs at least one non-padding target")
    return F.cross_entropy(logits.reshape(-1, logits.shape[-1]), target_ids.reshape(-1),
                           ignore_index=CharacterTokenizer.PAD_ID)
Download model.py
python -m tiny_gpt.inspect_model --save-initial checkpoints/initial.pt

This creates 413,440 random parameters and saves them. A parameter is an adjustable number. A checkpoint is a file containing those parameters so another process can load them later.

Run the untrained network once

Score 1+1=, then 1+2=. Compare the highest-scoring character. The network can produce scores, but its weights have not learned addition.

Runs on your device. First load is about 14 MB.

Random weights

Exact input: "1+1="

Score the prefix to inspect the next prediction.

Model and input details

Printable ASCII is passed verbatim. BOS is added once. All 98 tokens remain available. Context: 32 tokens. This exercise runs one forward pass and scores the next token. It does not generate an answer.

Random weights. Epoch 0. 1.73 MB of weights. The WASM runtime is downloaded on first use.

Technical details

This is a pre-layer-normalized causal Transformer with exact GELU, feed-forward width 512 and tied input/output embeddings. Attention uses softmax(QK transpose / sqrt(32)) with future positions masked. The source and export preserve the same layers.

Code and training steps