The concrete problem: predicting the next token
A language model has a simple training objective: given a sequence of tokens, predict the token that comes next. In a character-level model, a token is one character, so the task might be predicting e after th. The same objective applies when tokens represent words or subword pieces; only the vocabulary and sequence length change.
The engineering challenge is making each prediction depend on the useful parts of the preceding context without allowing the model to read the future. During training, the model receives many positions at once, but the prediction at position 10 must be based only on positions 1 through 10. This information boundary is what makes the model causal.
We can build a small causal transformer to make that mechanism tangible. It will read a text file, encode characters as integers, train on randomly selected context windows, and generate new text one character at a time. Because every component is visible, the relationship between the data, objective, architecture, and generated output is easier to inspect than in a large prebuilt system.
How it works
The model first maps each token ID to a learned vector. Token IDs themselves have no useful numeric meaning, so the embedding gives each character a trainable representation. It adds a learned positional vector because attention alone does not know whether a token appeared first or last. The resulting sequence is passed through transformer blocks, with one vector for each position.
Each attention head creates queries, keys, and values by multiplying the token representations by learned matrices. Query-key dot products measure which earlier positions are relevant to the current position; softmax turns those scores into weights; the weighted values are combined into a new representation. Different heads can learn different relationships, such as nearby character patterns or recurring delimiters.
Causal masking is the critical constraint. Before softmax, positions to the right of each token receive negative infinity, so their probability becomes zero. During training, every position can therefore predict its next token in parallel while still obeying the same information boundary used during generation. Without this mask, the model could use the answer it is supposed to predict and would not be learning the intended task.
A block adds the attention result back to its input, then applies a position-wise feed-forward network and another residual connection. Layer normalization stabilizes these transformations. The final linear layer converts each hidden vector into logits over the vocabulary, and cross-entropy compares those logits with the known next characters. The training target is the input sequence shifted one position to the left: each character supplies context for the character that follows it.
Here is a complete miniature implementation. Put training text in input.txt; the code uses only PyTorch and the Python standard library.
import torch
import torch.nn as nn
import torch.nn.functional as F
torch.manual_seed(1)
device = "cuda" if torch.cuda.is_available() else "cpu"
text = open("input.txt", encoding="utf-8").read()
chars = sorted(set(text))
stoi = {c: i for i, c in enumerate(chars)}
itos = {i: c for c, i in stoi.items()}
data = torch.tensor([stoi[c] for c in text], dtype=torch.long)
n, split = len(data), int(.9 * len(data))
train, valid = data[:split], data[split:]
block_size, batch_size = 128, 32
def batch(source):
starts = torch.randint(len(source) - block_size, (batch_size,))
x = torch.stack([source[i:i+block_size] for i in starts])
y = torch.stack([source[i+1:i+block_size+1] for i in starts])
return x.to(device), y.to(device)
class Head(nn.Module):
def __init__(self, dim, head_size):
super().__init__()
self.key, self.query, self.value = [nn.Linear(dim, head_size, False)
for _ in range(3)]
self.register_buffer("mask", torch.tril(torch.ones(block_size, block_size)))
def forward(self, x):
k, q, v = self.key(x), self.query(x), self.value(x)
scores = q @ k.transpose(-2, -1) * k.size(-1) ** -0.5
scores = scores.masked_fill(self.mask[:x.size(1), :x.size(1)] == 0, float("-inf"))
return F.softmax(scores, dim=-1) @ v
class Block(nn.Module):
def __init__(self, dim, heads):
super().__init__()
size = dim // heads
self.attn = nn.Linear(heads * size, dim)
self.heads = nn.ModuleList([Head(dim, size) for _ in range(heads)])
self.ff = nn.Sequential(nn.Linear(dim, 4 * dim), nn.ReLU(),
nn.Linear(4 * dim, dim))
self.n1, self.n2 = nn.LayerNorm(dim), nn.LayerNorm(dim)
def forward(self, x):
a = torch.cat([h(self.n1(x)) for h in self.heads], dim=-1)
x = x + self.attn(a)
return x + self.ff(self.n2(x))
class Model(nn.Module):
def __init__(self, vocab, dim=128, layers=4, heads=4):
super().__init__()
self.tok = nn.Embedding(vocab, dim)
self.pos = nn.Embedding(block_size, dim)
self.blocks = nn.Sequential(*[Block(dim, heads) for _ in range(layers)])
self.norm, self.lm_head = nn.LayerNorm(dim), nn.Linear(dim, vocab)
def forward(self, idx, targets=None):
t = idx.size(1)
x = self.tok(idx) + self.pos(torch.arange(t, device=device))
logits = self.lm_head(self.norm(self.blocks(x)))
loss = None if targets is None else F.cross_entropy(
logits.reshape(-1, logits.size(-1)), targets.reshape(-1))
return logits, loss
@torch.no_grad()
def generate(self, idx, steps):
for _ in range(steps):
context = idx[:, -block_size:]
logits, _ = self(context)
probs = F.softmax(logits[:, -1], dim=-1)
idx = torch.cat([idx, torch.multinomial(probs, 1)], dim=1)
return idx
model = Model(len(chars)).to(device)
opt = torch.optim.AdamW(model.parameters(), lr=3e-4)
for step in range(2000):
x, y = batch(train)
_, loss = model(x, y)
opt.zero_grad(set_to_none=True)
loss.backward()
opt.step()
if step % 200 == 0:
print(step, loss.item())
start = torch.zeros((1, 1), dtype=torch.long, device=device)
print("".join(itos[i] for i in model.generate(start, 500)[0].tolist()))
The batch function selects independent windows and creates the shifted targets. The generate method uses only the most recent block_size characters, computes probabilities for the final position, samples one character, and appends it before repeating. That loop is the training objective made operational: predict one next token, then use it as part of the next context.
What the evidence actually shows
This experiment demonstrates the architecture, not a competitive language model. On a small corpus, the model can learn spelling patterns, punctuation, local syntax, and recurring sequences; with enough capacity and training, it may mostly memorize the data. A convincing-looking sample can therefore reflect familiarity with the corpus rather than broad language understanding.
The useful evidence is in the loss and generated samples. Training loss should generally decline, while validation loss reveals whether the model is learning reusable structure or simply fitting the training text. The two views answer different questions: training loss describes optimization on seen examples, while validation loss tests behavior on held-out positions from the same data source.
If generated output becomes coherent only on familiar phrases, the model has learned distributional regularities but not robust generalization. Conversely, imperfect samples do not necessarily mean the mechanism is wrong; the corpus, number of updates, vocabulary, and context length all shape what the small model can learn. Samples are most useful when read alongside the loss.
The original transformer result established that attention-based sequence models could replace recurrent computation for major sequence tasks. Subsequent systems showed that increasing data, parameters, context, and training quality can produce much broader capabilities, but those results do not imply that a four-layer toy model will reason like a production model.
Using it
Start with a corpus large enough to fill many context windows, but small enough that iteration is quick. Plain text is useful because character tokenization removes vocabulary and tokenizer dependencies; it is also inefficient, since common words consume many positions. The character vocabulary is easy to inspect, but the model must spend several predictions spelling a word that a subword tokenizer might represent more compactly.
Watch the loss rather than treating a single sample as a metric. Add a validation batch and periodically evaluate it with model.eval() and torch.no_grad() in a real experiment. Save checkpoints containing model.state_dict() and the optimizer state so an interrupted run can resume. A checkpoint also lets you compare generations from the same trained state instead of accidentally comparing different training runs.
Change one variable at a time. Increase dim or the number of heads to test capacity, increase block_size to test longer dependencies, and adjust the learning rate before simply training longer. The attention matrix is also worth inspecting: it can show whether heads focus locally, attend to delimiters, or spread attention across the available context. Such inspection connects an abstract matrix operation to recognizable patterns in the text.
For reproducible comparisons, fix the random seed, keep the train/validation split stable, and record the parameter count, corpus, context length, and number of optimization steps. Generation temperature belongs to sampling, not training: lower values make output more predictable, while higher values explore lower-probability tokens. Keeping sampling settings fixed prevents a change in temperature from being mistaken for a change in model quality.
Limitations and failure modes
This implementation omits several production necessities. It uses character tokens, no dropout, no learning-rate schedule, no mixed precision, no distributed training, and no efficient attention kernel. Its feed-forward activation and normalization order are also deliberately simple. These omissions keep the implementation readable, but they also mean that training behavior and performance should not be generalized to larger systems.
Small datasets make memorization likely, and random context sampling can hide problems with document boundaries or duplicated text. A model may achieve a lower loss by exploiting formatting artifacts without learning the behavior you care about. Always inspect held-out examples and check whether generated text reproduces sensitive training material. Random windows can also mix contexts in ways that are convenient for training but do not represent the original document task.
The code assumes input.txt is longer than the context window and that its encoding is valid UTF-8. Larger models also make the quadratic attention cost visible: a sequence twice as long requires roughly four times as many pairwise attention scores per layer, before memory and optimization overheads. The context window therefore affects both what the model can connect and how expensive each forward pass becomes.
What to watch
The next useful step is not adding arbitrary layers; it is measuring what each change buys. Compare character and subword tokenization, visualize attention under controlled prompts, and test whether validation loss improves when the corpus contains longer-range structure. A controlled prompt makes it easier to distinguish a stable attention pattern from an appealing interpretation of one random sample.
From there, implement checkpointing and evaluation before scaling. A transformer becomes easier to reason about when its data pipeline, objective, masking rule, and sampling procedure are all small enough to inspect. Each change should have a corresponding observation: a loss curve, a held-out example, an attention visualization, or a generation under fixed sampling settings.
That discipline turns the miniature model into more than a toy. It provides a compact environment for asking concrete questions about context, capacity, optimization, and data quality, while keeping the code short enough that the answer can be traced back to a particular component.
Top comments (0)