Build a small transformer from scratch

A from-scratch character-level transformer trained on WikiText-2 reached a validation loss of 2.8717 using the best position-encoding scheme, according to a Stochastic Blog tutorial that isolates each design choice's cost and benefit. The model is a 64-dimensional, 2-layer, 4-head decoder with block size 64, trained for 60 steps on a 300,000-character slice of the training split, with every variant sharing the same seed, data, and parameter budget. After removing 2,054 duplicate training documents and 327 cross-split overlaps, the corpus holds 21,448 training, 2,361 validation, and 2,792 test documents, with a 1,013-token character vocabulary and a 0.000139 validation out-of-vocabulary rate.

Stochastic Blog https://stochastic.blog/tag/stochastic-blog/ Build a small transformer from scratch In the last three posts we built up the pieces of a language model: sequence modeling, tokenization, and the decoder-only setup that predicts the next token. Now we put the pieces together into an actual transformer. The through-line for this post is a single question we keep asking of every design choice: what does it cost, and what does it buy? A transformer is a stack of decisions about where to spend compute, and the small model we build here lets us measure each one. By the end, our best position scheme reaches a validation loss of 2.8717 on a character-level WikiText-2 decoder, and we will know exactly which choices earned that number. We work on WikiText-2, the raw release, about 13 MB of text and roughly two million word tokens under a CC BY-SA license. It is small enough to train a character-level decoder on a laptop CPU and still large enough to say something real. Everything below runs on a 64-dimensional, 2-layer, 4-head decoder with a block size of 64, and we train each variant for 60 steps on a 300,000-character slice of the training split. The numbers are small, but the comparisons are honest because every run shares the same seed, data, and parameter budget. What the data says Before building anything we look at the data. After dropping 2,054 duplicate training documents and 327 cross-split overlaps, we are left with 21,448 training documents, 2,361 validation documents, and 2,792 test documents. The character vocabulary is 1,013 tokens, and the out-of-vocabulary rate on validation is 0.000139, which is the reason a character tokenizer is the right call here: the vocabulary stays tiny and almost nothing is unseen. Figure 1 plots the size of each split alongside the most frequent characters, so we can see both the corpus shape and the character distribution in one view. The left panel shows the split sizes, and the right panel shows why the base rate matters. The most common character is a space, appearing at a base rate of 0.1892, so a model that always predicts a space is already at roughly 19 percent accuracy. That is the accuracy floor we have to beat, and it is a useful reference point for how much structure the model is capturing. Loss comparisons need a different baseline, though: to judge a cross-entropy number we would compare against the unigram predictor's loss, not against an accuracy rate. Figure 2 shows how document lengths are distributed, and the picture is a long-tailed one. The median document is 457 characters, but the distribution has a long tail out to 3,860 characters. A block size of 64 is a reasonable local context for this data, and it keeps every training step cheap. The word vocabulary is 76,603, which is another reason to stay at the character level for a from-scratch demo. Attention core The core operation of the transformer is scaled dot-product attention . Each token is projected into three vectors: a query , a key , and a value . The dot product between a query and a key decides how much of each value to mix in, and the result is softmax QKᵀ/√d V. The scaling by √d keeps the scores from saturating the softmax as the head dimension grows. The snippet below implements the formula with an optional causal mask. python def scaled dot product attention q, k, v, mask=None : d = q.size -1 scores = torch.matmul q, k.transpose -2, -1 / math.sqrt d similarity, scaled if mask is not None: scores = scores.masked fill ~mask, float '-inf' attn = F.softmax scores, dim=-1 return torch.matmul attn, v , attn The attention weights sum to 1 on each query row, which we verify directly: the row sums come out to 1.0, 0.99999994, 1.0, 1.0. Multi-head attention splits the model dimension across heads so each head attends to a different slice, and the heads are concatenated and projected back. In the unit test for this class, the model dimension is 16, so each of four heads sees four dimensions. Self-attention and cross-attention differ in where the queries and keys come from: in self-attention the queries, keys, and values all come from the same sequence and the sequence length is preserved, whereas in cross-attention a query sequence reads a different context sequence, which is how a decoder attends to an encoder. Causal masking is the constraint that makes a decoder a decoder: the mask is lower-triangular, so position i can only see positions up to and including i and never anything after it, which is what stops the model from peeking at the token it is supposed to predict. Where are we? Attention on its own is permutation invariant. Shuffle the tokens and the output shuffles with them, which means the model has no idea that word order matters. We need to inject position, and the four common schemes split into two families. The absolute schemes add a position vector to each token: sinusoidal positional encodings add fixed sine and cosine waves at different frequencies and cost no parameters, while learned positional embeddings add a trainable table with one vector per position. By contrast, the relative schemes inject distance directly into the attention computation: rotary position embeddings , or RoPE, rotate each query and key pair by an angle that depends on position so position enters through the dot product itself, and ALiBi adds a head-specific penalty proportional to the distance between query and key so far-away tokens are down-weighted.