Transformer block, explained with code

Transformer block

Also written: transformer blocks

One repeating unit of a transformer: attention, then a small neural network, each with a shortcut and normalization.

What it means

A transformer block is the unit that a transformer repeats. It does two jobs in turn. First comes attention, where each token gathers information from the others. Then a small neural network works on each token by itself. Each job has a layer norm in front and a shortcut around it, which keeps training stable. In a GPT, a mask stops each token from seeing the ones after it.

A complete block in PyTorch

import math
import torch
import torch.nn as nn

class Block(nn.Module):
    def __init__(self, d, heads):
        super().__init__()
        self.heads = heads
        self.ln1, self.ln2 = nn.LayerNorm(d), nn.LayerNorm(d)
        self.qkv, self.proj = nn.Linear(d, 3 * d), nn.Linear(d, d)
        self.mlp = nn.Sequential(nn.Linear(d, 4 * d), nn.GELU(), nn.Linear(4 * d, d))

    def forward(self, x):
        B, T, d = x.shape
        q, k, v = self.qkv(self.ln1(x)).split(d, dim=-1)
        q, k, v = (t.view(B, T, self.heads, d // self.heads).transpose(1, 2) for t in (q, k, v))
        scores = q @ k.transpose(-2, -1) / math.sqrt(d // self.heads)
        scores = scores.masked_fill(~torch.tril(torch.ones(T, T)).bool(), float("-inf"))
        x = x + self.proj((torch.softmax(scores, dim=-1) @ v).transpose(1, 2).reshape(B, T, d))
        return x + self.mlp(self.ln2(x))

torch.manual_seed(0)
block = Block(d=32, heads=4)
x = torch.randn(1, 5, 32)           # 1 sequence of 5 tokens, each a vector of 32 numbers
print("in: ", tuple(x.shape))
print("out:", tuple(block(x).shape))
print("parameters:", sum(p.numel() for p in block.parameters()))

What it prints

in:  (1, 5, 32)
out: (1, 5, 32)
parameters: 12704

The shape that comes out is the shape that went in, and that's the design's great trick: you can stack as many blocks as you like. A small model has a handful, and a large one has around a hundred. This block has 12,704 parameters, and most of a large model's parameters sit in blocks like it.

See it working

The transformer explainer runs a real tiny GPT made of blocks like this one, and build a GPT from scratch trains one yourself.

Try it now

Communicate, then compute, with a skip connection around each. Stack dozens of these and you have a GPT.

Other terms from the same lesson

decoder: A transformer that reads left to right and predicts the next token; each token can only see the ones before it. GPT-style chat models are decoders.

feed-forward network: The small neural network inside each transformer block. It works on each token on its own, after attention has mixed in information from the other tokens.

model card: The fact sheet published with a model: its size, what it was trained on, what it's for and its known limits.

SwiGLU: A feed-forward layer design used in Llama-style models, where one set of numbers acts as a gate that turns another set up or down.

transformer: The design behind today's language models: layers of attention plus small neural networks, repeated many times.

Never trained an AI?

Teach a small one to write, right in your browser. It takes about 20 seconds and needs no account.

Teach it to write