Build the network
We've turned the text into numbers. Now we need a program that uses those numbers to predict the next character. The part that makes this prediction is a neural network.
The network uses stored numbers called weights in its calculations. During training, we adjust these weights using examples. Before we train it, let's see how one weight works.
What is a weight?
A neuron is a small calculation inside a neural network. In this example, it multiplies each input by a weight, adds the results and turns any negative sum into zero.
Move the slider to change the first weight. Watch the result on the right and the output below it. Try a negative weight too.
A neural network connects many calculations like this. One neuron's output can become another's input.
Why add zero? Why change negative numbers?
The extra number added to the sum is called a bias. It starts at zero here and can change during training, just like a weight.
The rule that turns negative sums into zero is called ReLU. It's an example of a nonlinear function. Without nonlinear functions, these stacked linear calculations could be replaced by one linear calculation.
Our model uses a smoother function called GELU in its feed-forward networks. The inputs in this demo are made-up internal values. They aren't the digits from an addition question.
Try our network
Score the next character for 1+1=, then 1+2=. Compare the results from the untrained network.
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.
Optional code
Build this network in Python
Each token's embedding is a list of 128 numbers. The embedding table stores one list per token ID. The network carries these lists in an array shaped [batch, positions, 128].
Each layer is a Python class based on nn.Module. Its __init__ method creates the weights. forward defines the calculation.
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)
python -m tiny_gpt.inspect_model --save-initial checkpoints/initial.ptThis creates and saves 413,440 parameters. Most weights start with small random values. Biases start at zero and normalization scales start at one. The saved file is called a checkpoint. It lets another program load the same model later.
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