Capstone: a tiny attention block

The block in one diagram

A transformer block is two sublayers, each with a residual connection around it. A residual connection adds a sublayer’s output back to its input. The first sublayer normalises the stream, runs multi-head self-attention and adds the result back. The second normalises the updated stream, runs the feed-forward network and adds that result back. The output has the same shape as the input. The block can therefore be stacked without any reshaping.

The capstone builds exactly that. It uses four-dimensional tokens, two heads of width two, a hidden layer of width eight and a causal mask. The numbers are small enough to print and check by hand. The structure is the same one that appears in a model with a hundred layers and a width of thousands. Only the sizes and a few details change.

The two residual adds are not decoration. The first keeps the token’s identity while attention writes context onto it. The second keeps the context while the feed-forward block transforms it. Remove either and the block becomes a replacement, not an update. A replacement is much harder to stack and much easier to destabilise. The normalisation before each branch keeps the update at a workable scale without touching the stream itself.

What to verify, not only run

A block that runs is not necessarily correct. Three checks catch most mistakes. The output shape must equal the input shape. A causal model must give the same first-position output when only a later token changes. The residual stream should also preserve a path for the input. The output should therefore not be at a wildly different scale from the input, unless the updates are large.

The first check catches a projection with the wrong width. The second catches a mask applied after the softmax, or not at all. The third catches a missing residual or a normalisation on the wrong axis. None of these checks needs a trained model. They belong in a test, not in a training run.

The same block, stacked

from torch import nn

layer = nn.TransformerEncoderLayer(
    d_model=64,
    nhead=4,
    dim_feedforward=256,
    batch_first=True,
    norm_first=True,
)
encoder = nn.TransformerEncoder(layer, num_layers=6)
out = encoder(x)
print(out.shape)  # (batch, tokens, 64)

Reading the output

The block returns a vector per token, not a prediction. To turn it into a language model, apply a final layer norm. A linear head then maps each vector to vocabulary logits. The last position’s logits are the distribution for the next token. During training, every position’s logits are compared with the shifted targets.

The block is also not a model until its weights are trained. Random projections produce a valid tensor and no useful language. The capstone shows the plumbing: shapes, masks, residuals and normalisation. The learning is a separate loop. It adjusts the numbers while the plumbing stays fixed.

The final norm matters more than it looks. The residual stream can grow with depth, and the output head expects a vector at a consistent scale. A missing final norm produces logits that are too large. The softmax then becomes nearly one-hot. The model looks confident before it has learned anything useful. The head itself is usually a single matrix, sometimes tied to the input embeddings. It is the only part that knows the vocabulary size.

Where the real model differs

Production models stack dozens of these blocks. They widen the stream, replace learned positions with RoPE and often replace the feed-forward with a gated variant. They use grouped-query attention to shrink the inference cache. They use mixed precision to fit the weights, and dropout or weight decay to regularise. Those are changes to the parts, not to the skeleton.

The skeleton is what makes the rest legible. Trace a token through one block and a stack becomes repetition. A training run is that repetition plus a loss and an optimiser. The details matter. They are easier to place when the small version is already in your head.

The same skeleton explains why architecture changes are often small in code and large in effect. Swapping the activation, moving the norm or changing the head count touches a few lines. It changes the training landscape and the inference cost. The capstone is a place to make those changes and watch the shapes and the causality test. A training run would hide the mistake in a loss curve.