Jekyll 2026-11-11T10:59:38+00:00 https://blog.korolevalab.com/feed.xml Eva Koroleva Eva Koroleva, ML engineer and researcher in NLP, computer vision and generative models: projects, a Transformers course, and posts on ML and cognitive science. Eva Koroleva https://blog.korolevalab.com/ The Vision Transformer reads images as sequences of patches 2026-11-09T20:00:00+00:00 2026-11-10T00:00:00+00:00 https://blog.korolevalab.com/blog/2026/transformer-architectures-11-vit Transformer architectures: from BERT to diffusion · lesson 11 of 15 · all lessons · run on Kaggle · open in Colab · code on GitHub

For eight years, from AlexNet in 2012 to 2020, every serious image classifier was a convolutional network. A convolution has two beliefs about images built into it: that nearby pixels belong together, and that a pattern means the same thing wherever it appears. Those beliefs are true, and they are why a CNN learns from a few thousand photos.

In October 2020 a team at Google Brain dropped both. They took the Transformer encoder of lesson 1, BERT’s encoder almost unchanged, cut an image into a grid of 16×16-pixel squares, and fed the squares to it as if they were words.

The paper’s title was the whole idea: An image is worth 16x16 words. It worked, with a condition: the model needed to see far more images than a CNN, 303 million of them in Google’s internal JFT-300M collection, before it caught up and then overtook.

Today that same encoder, the Vision Transformer or ViT, is the eye of every model that reads images, from DINOv2 to the vision-language models of lesson 13.

One thing to settle before we start, because the phrase “images as words” invites it: the ViT doesn’t turn an image into text. It turns an image into a sequence of vectors, the same kind of object an embedding layer makes from text, and the Transformer reads that sequence with exactly the machinery it uses for words.

A vision-language model (lesson 13) goes one step further and places these vectors next to the vectors of text tokens in a language model’s input. Generating images from text is the opposite direction and has its own architecture, lesson 15’s diffusion Transformer.

The Vision Transformer in one line. A 224 by 224 photo is cut into a 14 by 14 grid of 16-pixel squares; each square is flattened into 768 numbers and projected to a vector of 768; a learned class token is put in front and a learned position vector is added to each of the 197 vectors; the sequence goes through twelve encoder blocks, and the class token's output vector goes to a linear classifier that names the picture The Vision Transformer in one line. A 224 by 224 photo is cut into a 14 by 14 grid of 16-pixel squares; each square is flattened into 768 numbers and projected to a vector of 768; a learned class token is put in front and a learned position vector is added to each of the 197 vectors; the sequence goes through twelve encoder blocks, and the class token's output vector goes to a linear classifier that names the picture

What you’ll do

  1. Walk one photo from pixels to the sequence of 197 vectors the encoder reads, one step at a time, with every tensor’s shape printed and every step drawn.
  2. Write the ViT in PyTorch, load Google’s ViT-B/16 into it and match Hugging Face.
  3. Open the trained model: what its first layer looks for in a patch, what its positions learned, how far its heads look at each depth, which parts of a photo decide the answer.
  4. Run it at a higher resolution by stretching its position embeddings, lesson 9’s trick in two dimensions.
  5. Watch DINO, a ViT trained without a single label, find the objects in photos with its attention.
  6. On satellite images of Europe, race a pretrained ViT, a ViT and a CNN trained from scratch, and a probe on DINOv2’s features, from a hundred labelled images to sixteen thousand.

What you need: a GPU for the last section. On Kaggle, Settings → Accelerator → GPU; on Colab, Runtime → Change runtime type → T4 GPU. The notebook downloads ViT-B/16 (346 MB) twice, in two versions, two small self-supervised ViTs (90 MB each), five public-domain photos and the EuroSAT dataset (90 MB).

Setup

One small helper package, tfs, draws the course’s plots and token chips and knows where the notebook runs (Kaggle, Colab or your own machine). Everything the lesson teaches is written out in the notebook itself.

import importlib.util
import subprocess
import sys

TFS = "git+https://github.com/xmarva/transformers-from-scratch@main"
if importlib.util.find_spec("tfs") is None:
    subprocess.check_call([sys.executable, "-m", "pip", "install", "-q", TFS])

from tfs import env
from tfs.viz import table, token_chips, use_course_style

env.quiet()  # no progress bars or environment warnings in the outputs
SEED = 42
env.seed_everything(SEED)
COLORS = use_course_style()
WORK_DIR = env.work_dir()

import io
import json
import math
import random
import time
import urllib.request

import matplotlib.pyplot as plt
from matplotlib.colors import LinearSegmentedColormap, to_rgb
import numpy as np
import pandas as pd
import torch
import torch.nn as nn
import torch.nn.functional as F
from huggingface_hub import hf_hub_download
from PIL import Image
from safetensors.torch import load_file

SMOKE = env.smoke()  # CI runs the lesson on a small slice; you run it in full
device = "cuda" if torch.cuda.is_available() else "cpu"
if device == "cuda":
    print("GPU:", torch.cuda.get_device_name(0))

GPU: Tesla T4

1. Five photographs

Five photographs from the collections of the US National Park Service, NASA, the US Geological Survey, NOAA and the US Fish and Wildlife Service: works of the US federal government, in the public domain.

A brown bear catching a salmon at Brooks Falls in Alaska, the space shuttle Discovery lifting off, a fountain of lava on Kīlauea, an orca breaching off the coast of Washington, a rufous hummingbird in Wyoming.

PHOTOS = ["bear", "shuttle", "lava", "orca", "hummingbird"]
# 960-pixel copies of the Wikimedia Commons files listed in the References, kept in
# the course's repository because Commons limits how often one address may download
BASE = (
    "https://raw.githubusercontent.com/"
    "xmarva/transformers-from-scratch/main/assets/photos/a11"
)


def fetch_photo(name):
    data = urllib.request.urlopen(f"{BASE}/{name}.jpg").read()
    return Image.open(io.BytesIO(data)).convert("RGB")


photos = {}
for name in PHOTOS:
    photos[name] = fetch_photo(name)

fig, axes = plt.subplots(1, len(photos), figsize=(12, 2.8))
for ax, (key, photo) in zip(axes, photos.items()):
    ax.imshow(photo)
    ax.set_title(f"{key}, {photo.width}×{photo.height}", fontsize=10)
    ax.axis("off")
plt.show()

Output of the code above

2. From pixels to a sequence, step by step

This is the section the rest of the lesson stands on. We take the bear and follow it to the exact tensor the encoder reads, printing the shape at every step. The numbers come from the model we load in section 3, Google’s ViT-B/16: patches of 16 pixels, vectors of 768 numbers.

Step 1: a fixed size and a fixed scale. The model was trained on 224×224 images, so every photo is resized to that, aspect ratio or not.

Each pixel’s three channels, red, green and blue, arrive as integers from 0 to 255; they are divided by 255 and then normalised, shifted and scaled so that each channel sits roughly between −1 and 1. This checkpoint uses a mean of 0.5 and a standard deviation of 0.5 for every channel:

\[x_{c,i,j} = \frac{\text{pixel}_{c,i,j} / 255 - 0.5}{0.5}.\] 
SIZE, PATCH, MEAN, STD = 224, 16, 0.5, 0.5


def to_tensor(photo, size=SIZE, mean=MEAN, std=STD):
    small = photo.resize((size, size), Image.BILINEAR)
    pixels = torch.tensor(np.asarray(small), dtype=torch.float32)  # (height, width, 3)
    pixels = (
        pixels.permute(2, 0, 1) / 255
    )  # channels first: (3, height, width), from 0 to 1
    mean = torch.tensor(mean).reshape(-1, 1, 1)
    std = torch.tensor(std).reshape(-1, 1, 1)
    return (pixels - mean) / std


x = to_tensor(photos["bear"])
print("one photo as a tensor:", tuple(x.shape))
print(f"values from {x.min():.2f} to {x.max():.2f}")

fig, axes = plt.subplots(1, 3, figsize=(11, 3.2), gridspec_kw={"width_ratios": [1, 1, 1.4]})
axes[0].imshow(photos["bear"])
axes[0].set_title(f"original, {photos['bear'].width}×{photos['bear'].height}", fontsize=10)
axes[1].imshow((x.permute(1, 2, 0) * STD + MEAN).clamp(0, 1))
axes[1].set_title("resized to 224×224", fontsize=10)
for channel, color in zip(range(3), [COLORS["dec"], COLORS["enc"], COLORS["att"]]):
    axes[2].hist(
        x[channel].flatten(), bins=60, color=color, alpha=0.6, label="RGB"[channel]
    )
axes[2].set_title("the 3 × 224 × 224 numbers after normalising", fontsize=10)
axes[2].legend(frameon=False)
axes[0].axis("off")
axes[1].axis("off")
plt.show()

one photo as a tensor: (3, 224, 224)
values from -0.96 to 1.00

resized to 224×224

The bear became a little wider, because the photo wasn’t square; the model was trained on such squashed squares and doesn’t mind. The histogram shows what the model actually receives: 150,528 numbers, mostly in the bright end, because the photo is mostly white water.

Step 2: cut into patches. A 224×224 image holds 14 × 14 = 196 squares of 16×16 pixels. Cutting is pure reindexing, no arithmetic: split the height into 14 rows of 16 and the width into 14 columns of 16, move the row and column indices to the front, and flatten each square.

def patchify(x, p=PATCH):
    channels, height, width = x.shape
    x = x.reshape(channels, height // p, p, width // p, p)  # (3, 14, 16, 14, 16)
    x = x.permute(
        1, 3, 0, 2, 4
    )  # (14, 14, 3, 16, 16): grid row, grid column, then the patch
    return x.reshape((height // p) * (width // p), channels * p * p)  # (196, 768)


patches = patchify(x)
print("patches:", tuple(patches.shape))
print("numbers in one patch: 3 × 16 × 16 =", 3 * PATCH * PATCH)

patches: (196, 768)
numbers in one patch: 3 × 16 × 16 = 768

The order of the squares is the order of reading a page: left to right along the top row, then the next row. Here is the grid on the photo, and the first two rows of squares laid out as the sequence the model will see:

grid = SIZE // PATCH
shown = (x.permute(1, 2, 0) * STD + MEAN).clamp(0, 1)
fig = plt.figure(figsize=(12, 4.2))
ax = fig.add_axes([0.0, 0.0, 0.32, 1.0])
ax.imshow(shown)
for k in range(1, grid):
    ax.axhline(k * PATCH - 0.5, color="white", linewidth=0.6)
    ax.axvline(k * PATCH - 0.5, color="white", linewidth=0.6)
red = COLORS["dec"]
first_rows = plt.Rectangle((-0.5, -0.5), SIZE, 2 * PATCH, fill=False, color=red, lw=2)
ax.add_patch(first_rows)
ax.set_title("a 14 × 14 grid of 16-pixel squares; the first two rows in red", fontsize=10)
ax.axis("off")
for k in range(2 * grid):
    row, col = divmod(k, grid)
    tile = shown[row * PATCH : (row + 1) * PATCH, col * PATCH : (col + 1) * PATCH]
    tile_ax = fig.add_axes(
        [0.35 + (k % grid) * 0.046, 0.62 - (k // grid) * 0.3, 0.042, 0.25]
    )
    tile_ax.imshow(tile)
    tile_ax.set_xticks([])
    tile_ax.set_yticks([])
    tile_ax.set_title(str(k), fontsize=7)
caption = "patches 0 to 27 of 196, in reading order; the sequence goes on row by row"
fig.text(0.35, 0.05, caption, fontsize=10, color=COLORS["muted"])
plt.show()

a 14 × 14 grid of 16-pixel squares; the first two rows in red

PyTorch has the same cut built in as unfold, which slides a window over an image and stacks what it sees. With a stride equal to the window, the windows don’t overlap and it returns our patches, in our order:

unfolded = F.unfold(x[None], kernel_size=PATCH, stride=PATCH)[0].T  # (196, 768)
print("unfold against reshape, max difference:", (unfolded - patches).abs().max().item())

unfold against reshape, max difference: 0.0

Step 3: one patch is 768 numbers. A patch is a small picture, 16×16 pixels in three channels, and the model sees it as one long vector: the 256 red values, then the 256 green, then the 256 blue.

Here is patch 50, the one with the bear’s eye, and its vector drawn as a strip, every number one thin bar.

k = 50  # row 3, column 8: the bear's eye
row, col = divmod(k, grid)
tile = shown[row * PATCH : (row + 1) * PATCH, col * PATCH : (col + 1) * PATCH]
fig = plt.figure(figsize=(12, 2.6))
ax = fig.add_axes([0.0, 0.05, 0.16, 0.85])
ax.imshow(tile, interpolation="nearest")
ax.set_title(f"patch {k}: row {row}, column {col}", fontsize=10)
ax.axis("off")
ax = fig.add_axes([0.2, 0.25, 0.79, 0.5])
greys = LinearSegmentedColormap.from_list("greys", [COLORS["ink"], COLORS["card"]])
ax.imshow(patches[k][None], aspect="auto", cmap=greys, interpolation="nearest")
for edge in [256, 512]:
    ax.axvline(edge - 0.5, color=COLORS["dec"], linewidth=1.5)
ax.set_yticks([])
ax.set_xticks([128, 384, 640], ["256 red values", "256 green values", "256 blue values"])
ax.set_title("the same patch as the 768-number vector the model reads", fontsize=10)
plt.show()

the same patch as the 768-number vector the model reads

Step 4: the patch embedding. Text enters a Transformer through an embedding table: token number 1,723 picks row 1,723 of a matrix.

A patch has no number; it is a vector of 768 pixels. So the ViT multiplies it by a learned matrix \(E\) and adds a bias:

\[z_i = E\,p_i + b, \qquad E \in \mathbb{R}^{D \times 768}.\]

The two are closer than they look. Looking up row \(t\) of a table is multiplying the table by a one-hot vector, a vector of zeros with a single 1 in place \(t\); the ViT multiplies the same kind of matrix by a vector of pixels instead of a one-hot one.

How text and images become vectors, side by side. On the left, a token id picks one row of an embedding table, which is the same as multiplying the table by a one-hot vector. On the right, a patch of 16 by 16 pixels in three channels is flattened into 768 numbers and multiplied by a learned matrix, plus a bias, giving a vector of the same width as a word's embedding. Both produce one vector per token; after that the Transformer can't tell where the vectors came from How text and images become vectors, side by side. On the left, a token id picks one row of an embedding table, which is the same as multiplying the table by a one-hot vector. On the right, a patch of 16 by 16 pixels in three channels is flattened into 768 numbers and multiplied by a learned matrix, plus a bias, giving a vector of the same width as a word's embedding. Both produce one vector per token; after that the Transformer can't tell where the vectors came from

In the released code the matrix is stored as a convolution with 768 filters of size 16×16 and a stride of 16: sliding a 16×16 filter in steps of 16 is the same as multiplying each patch by the filter, flattened. We check it with the real weights.

from transformers import ViTConfig, ViTForImageClassification

REPO = "google/vit-base-patch16-224"
if SMOKE:  # a small model with random weights stands in for Google's
    small = ViTConfig(hidden_size=64, num_hidden_layers=2, num_attention_heads=4)
    small.intermediate_size, small.num_labels = 128, 1000
    torch.manual_seed(0)
    reference = ViTForImageClassification(small).eval()
    vit_config = small.to_dict()
    reference.save_pretrained(WORK_DIR / "vit-small")  # saved under the checkpoints' names
    hf_weights = load_file(str(WORK_DIR / "vit-small" / "model.safetensors"))
else:
    reference = ViTForImageClassification.from_pretrained(REPO).eval()
    vit_config = json.load(open(hf_hub_download(REPO, "config.json")))
    hf_weights = load_file(hf_hub_download(REPO, "model.safetensors"))
conv_weight = hf_weights[
    "vit.embeddings.patch_embeddings.projection.weight"
]  # (768, 3, 16, 16)
conv_bias = hf_weights["vit.embeddings.patch_embeddings.projection.bias"]
print("the convolution's weight:", tuple(conv_weight.shape))

E = conv_weight.reshape(
    conv_weight.shape[0], -1
)  # (768, 768): one flattened filter per row
as_matrix = patches @ E.T + conv_bias  # (196, 768)
as_conv = F.conv2d(x[None], conv_weight, conv_bias, stride=PATCH)  # (1, 768, 14, 14)
as_conv = as_conv.flatten(2)[0].T  # (196, 768)
print("patch vectors:", tuple(as_matrix.shape))
print(
    "matrix against convolution, max difference:",
    f"{(as_matrix - as_conv).abs().max():.1e}",
)

the convolution's weight: (768, 3, 16, 16)
patch vectors: (196, 768)
matrix against convolution, max difference: 9.5e-06

The convolution is only a fast way to write the matrix product: the ViT has no convolutional layers in the sense of a CNN, because the filters never overlap and never slide past what they’ve already seen.

Step 5: the class token and the positions. Two learned additions turn the 196 vectors into the sequence the encoder reads.

In front goes a class token, one vector the model learned, the same for every image: it carries no pixels and gathers the whole image through attention, layer by layer, like BERT’s [CLS] in lesson 1.

Its output is what the classifier reads. Then, to every one of the 197 vectors, a learned position embedding is added, one per place in the grid, because attention by itself has no idea of order.

cls_token = hf_weights["vit.embeddings.cls_token"]  # (1, 1, 768)
positions = hf_weights["vit.embeddings.position_embeddings"]  # (1, 197, 768)
sequence = torch.cat([cls_token[0], as_matrix], dim=0)  # (197, 768)
sequence = sequence + positions[0]
print("class token:", tuple(cls_token.shape), " positions:", tuple(positions.shape))
print("the sequence the encoder reads:", tuple(sequence.shape))

# negative numbers in teal, positive in rose, zero white
ends = [COLORS["enc"], COLORS["card"], COLORS["dec"]]
signed = LinearSegmentedColormap.from_list("signed", ends)
fig, ax = plt.subplots(figsize=(11, 2.6))
ax.imshow(
    sequence.T[:96], aspect="auto", cmap=signed, vmin=-3, vmax=3, interpolation="nearest"
)
ax.set_xlabel("position in the sequence: 0 is the class token, 1 to 196 the patches")
ax.set_ylabel("first 96 of 768")
ax.set_title("The bear, as the encoder sees it: 197 columns of 768 numbers", fontsize=10)
plt.show()

class token: (1, 1, 768)  positions: (1, 197, 768)
the sequence the encoder reads: (197, 768)

The bear, as the encoder sees it: 197 columns of 768 numbers

That is the whole conversion. From here on, a ViT is a Transformer encoder, the one from lesson 1: twelve blocks of attention and an MLP, then a final layer norm, and a linear layer on the class token’s output vector.

Are the positions really needed? Attention compares every vector with every other one and sums the results, so without positions the encoder can’t tell a patch in the corner from the same patch in the centre.

A test makes it concrete: shuffle the 196 patches, as if the photo were cut up and the pieces thrown in a bag. We do it twice, with the positions the model learned and with the positions removed. The encode function used here is the encoder of the next section.

3. The ViT in PyTorch

The encoder block puts the layer norm before attention and before the MLP, as GPT-2 and every model since do, not after them as BERT did; the MLP widens to 3,072 with a GELU, as in BERT. The attention keeps its weights on demand, for the pictures of section 4.

class Attention(nn.Module):
    def __init__(self, d, heads):
        super().__init__()
        self.heads, self.head_dim = heads, d // heads
        self.query = nn.Linear(d, d)
        self.key = nn.Linear(d, d)
        self.value = nn.Linear(d, d)
        self.out = nn.Linear(d, d)
        self.keep = False
        self.weights = None

    def forward(self, x):
        b, n, d = x.shape
        q = self.query(x).view(b, n, self.heads, self.head_dim).transpose(1, 2)
        k = self.key(x).view(b, n, self.heads, self.head_dim).transpose(1, 2)
        v = self.value(x).view(b, n, self.heads, self.head_dim).transpose(1, 2)
        if self.keep:  # the slow way, to see the weights
            weights = (q @ k.transpose(-1, -2) / math.sqrt(self.head_dim)).softmax(-1)
            self.weights = weights.detach()
            out = weights @ v
        else:
            out = F.scaled_dot_product_attention(q, k, v)
        return self.out(out.transpose(1, 2).reshape(b, n, d))


class Block(nn.Module):
    def __init__(self, d, heads, hidden, eps):
        super().__init__()
        self.norm1 = nn.LayerNorm(d, eps=eps)
        self.attn = Attention(d, heads)
        self.norm2 = nn.LayerNorm(d, eps=eps)
        self.fc1 = nn.Linear(d, hidden)
        self.fc2 = nn.Linear(hidden, d)

    def forward(self, x):
        x = x + self.attn(self.norm1(x))
        return x + self.fc2(F.gelu(self.fc1(self.norm2(x))))

The model is the five steps of section 2 followed by the blocks. tokens builds the sequence; encode runs the blocks on any sequence, which is what the shuffling test needs.

class ViT(nn.Module):
    def __init__(self, c, classes, image_size=None):
        super().__init__()
        d, p = c["hidden_size"], c["patch_size"]
        grid = (image_size or c["image_size"]) // p
        eps = c.get("layer_norm_eps", 1e-6)
        self.patch_embed = nn.Conv2d(3, d, kernel_size=p, stride=p)
        self.cls = nn.Parameter(torch.zeros(1, 1, d))
        self.pos = nn.Parameter(torch.zeros(1, grid * grid + 1, d))
        heads, hidden = c["num_attention_heads"], c["intermediate_size"]
        self.blocks = nn.ModuleList(
            [Block(d, heads, hidden, eps) for _ in range(c["num_hidden_layers"])]
        )
        self.norm = nn.LayerNorm(d, eps=eps)
        self.head = nn.Linear(d, classes)

    def tokens(self, pixels):
        x = self.patch_embed(pixels).flatten(2).transpose(1, 2)  # (batch, 196, d)
        x = torch.cat([self.cls.expand(x.shape[0], -1, -1), x], dim=1)  # (batch, 197, d)
        return x + self.pos

    def encode(self, x):
        for block in self.blocks:
            x = block(x)
        return self.norm(x)

    def forward(self, pixels):
        return self.head(self.encode(self.tokens(pixels))[:, 0])

Google’s names differ from ours; one function renames them. The same function serves every ViT checkpoint in this lesson.

def rename(state):
    short = {
        "embeddings.patch_embeddings.projection": "patch_embed",
        "embeddings.cls_token": "cls",
        "embeddings.position_embeddings": "pos",
        "layernorm": "norm",
        "classifier": "head",
    }
    inside = {
        "attention.attention.": "attn.",
        "attention.output.dense": "attn.out",
        "layernorm_before": "norm1",
        "layernorm_after": "norm2",
        "intermediate.dense": "fc1",
        "output.dense": "fc2",
    }
    out = {}
    for key, value in state.items():
        key = key.removeprefix("vit.")
        if key.startswith("pooler."):
            continue  # a layer this checkpoint's classifier never used
        if key.startswith("encoder.layer."):
            number, rest = key.removeprefix("encoder.layer.").split(".", 1)
            for theirs, ours in inside.items():
                rest = rest.replace(theirs, ours)
            out[f"blocks.{number}.{rest}"] = value
            continue
        for theirs, ours in short.items():
            if key.startswith(theirs):
                key = ours + key[len(theirs) :]
        out[key] = value
    return out

The check against Hugging Face, whose model was loaded in section 2, on all five photos at once.

labels = {int(k): v.split(",")[0] for k, v in vit_config["id2label"].items()}
vit = ViT(vit_config, len(labels))
vit.load_state_dict(rename(hf_weights), strict=True)
vit.eval()

batch = torch.stack([to_tensor(photo) for photo in photos.values()])
with torch.no_grad():
    ours = vit(batch)
    theirs = reference(pixel_values=batch).logits
print("batch:", tuple(batch.shape), " logits:", tuple(ours.shape))
print("max difference from Hugging Face:", f"{(ours - theirs).abs().max().item():.1e}")

batch: (5, 3, 224, 224)  logits: (5, 1000)
max difference from Hugging Face: 0.0e+00

Exact to the last bit: same operations, same order. Now the shuffling test from section 2, on the five photos.

torch.manual_seed(SEED)
order = torch.randperm(grid * grid)


def predict(with_positions, shuffled):
    x = vit.patch_embed(batch).flatten(2).transpose(1, 2)
    if shuffled:
        x = x[:, order]
    x = torch.cat([vit.cls.expand(len(batch), -1, -1), x], dim=1)
    if with_positions:
        x = x + vit.pos
    with torch.no_grad():
        probs = vit.head(vit.encode(x)[:, 0]).softmax(-1)
    return probs


rows = []
for with_positions in [True, False]:
    for shuffled in [False, True]:
        probs = predict(with_positions, shuffled)
        row = {"positions": "learned" if with_positions else "removed"}
        row["patches"] = "shuffled" if shuffled else "in order"
        for name, p in zip(photos, probs):
            row[name] = f"{labels[p.argmax().item()]} {p.max().item():.0%}"
        rows.append(row)
table(pd.DataFrame(rows))
positionspatchesbearshuttlelavaorcahummingbird
learnedin orderbrown bear 100%space shuttle 100%volcano 100%killer whale 100%hummingbird 100%
learnedshuffledbrown bear 83%projectile 58%volcano 100%kite 26%hummingbird 28%
removedin ordergolden retriever 50%honeycomb 65%volcano 43%black stork 14%bulbul 18%
removedshuffledgolden retriever 50%honeycomb 65%volcano 43%black stork 14%bulbul 18%

Two things to read. Without positions, shuffled and ordered patches give identical answers, to the last digit: with nothing to say where a patch sits, the encoder treats the image as a bag of squares, and a bag has no order. That model also gets most photos wrong, because it was trained to rely on positions it no longer has.

With positions, shuffling hurts some photos and not others: the bear’s fur and the lava’s red are recognisable in any order, while a shuttle cut into pieces looks like any projectile, and a shuffled orca becomes a kite, the bird of prey.

ViTs are surprisingly robust to such shuffling, a finding Naseer and colleagues reported in 2021, and the reason is in the table: much of what identifies an object lives inside the patches themselves.

The model’s own answers on the five photos, top five each:

probs = ours.softmax(-1)
fig, axes = plt.subplots(
    2, len(photos), figsize=(13, 4.8), gridspec_kw={"height_ratios": [1.2, 1]}
)
for k, photo in enumerate(photos.values()):
    axes[0, k].imshow(photo.resize((SIZE, SIZE)))
    axes[0, k].axis("off")
    top = probs[k].topk(5)
    names = [labels[i.item()][:18] for i in top.indices][::-1]
    axes[1, k].barh(names, top.values.flip(0).numpy(), color=COLORS["att"])
    axes[1, k].set_xlim(0, 1)
    axes[1, k].tick_params(labelsize=8)
plt.tight_layout()
plt.show()

Output of the code above

4. What the model learned

The first layer. The patch embedding’s 768 filters are 768 small pictures, each a pattern the model compares a patch against.

Their principal components, the directions along which the filters vary most, show what kind of patterns those are. The ViT paper drew the same picture for its large model.

filters = vit.patch_embed.weight.detach().reshape(vit.patch_embed.weight.shape[0], -1)
centered = filters - filters.mean(dim=0)
_, _, directions = torch.linalg.svd(centered, full_matrices=False)
n_show = min(28, directions.shape[0])
fig, axes = plt.subplots(4, 7, figsize=(8.4, 4.8))
for k, ax in enumerate(axes.flat):
    ax.axis("off")
    if k >= n_show:
        continue
    picture = directions[k].reshape(3, PATCH, PATCH).permute(1, 2, 0)
    picture = (picture - picture.min()) / (picture.max() - picture.min())
    ax.imshow(picture, interpolation="nearest")
fig.suptitle("The 28 main directions of ViT-B/16's patch filters", fontsize=11)
plt.show()

Output of the code above

Blobs of colour, colour-opponent pairs, then edges and spots at every orientation, then finer and finer textures: much the same basis as the first layer of a CNN, which also learns edge and colour detectors. Nobody built them in; they are what is useful for describing a 16×16 square of a natural image.

The positions. Each of the 196 position embeddings is a vector of 768 numbers that the model learned with no hint that the patches form a grid. Compare each position with every other one by cosine similarity, and draw the result for each place in the grid as a small map:

pos = vit.pos[0, 1:].detach()
similarity = F.cosine_similarity(pos[:, None], pos[None], dim=-1)
similarity = similarity.reshape(grid, grid, grid, grid)
step = 2  # every second row and column, so that the maps stay readable
picked = list(range(0, grid, step))
fig, axes = plt.subplots(len(picked), len(picked), figsize=(7.6, 7.6))
for a, i in enumerate(picked):
    for b, j in enumerate(picked):
        axes[a, b].imshow(
            similarity[i, j], cmap=signed, vmin=-1, vmax=1, interpolation="nearest"
        )
        axes[a, b].set_xticks([])
        axes[a, b].set_yticks([])
fig.suptitle(
    "Each map: one grid position against all 196, by cosine similarity", fontsize=11
)
plt.show()

Output of the code above

Every map is strongest exactly where its position is, with a halo around it and bands of rose along its own row and column.

The model rediscovered the two-dimensional grid from a one-dimensional list of 196 slots, simply because knowing which patches are neighbours, and which share a row or a column, helps it classify. The ViT paper tried position embeddings designed for two dimensions and found no gain; this picture is why.

How far the heads look. For each head, the mean attention distance is the distance in pixels between a patch and the patches it attends to, averaged with the attention weights. A head that looks at its neighbours has a small distance, one that looks across the whole image a large one.

for block in vit.blocks:
    block.attn.keep = True
with torch.no_grad():
    vit(batch)
centres = (
    torch.stack(
        torch.meshgrid(torch.arange(grid), torch.arange(grid), indexing="ij"), dim=-1
    )
    .reshape(-1, 2)
    .float()
    * PATCH
)
pixels_apart = torch.cdist(centres, centres)  # (196, 196)
distance = []
for block in vit.blocks:
    w = block.attn.weights[:, :, 1:, 1:]  # patches to patches, without the class token
    w = w / w.sum(-1, keepdim=True)
    distance.append((w * pixels_apart).sum(-1).mean(dim=(0, 2)))  # one number per head
distance = torch.stack(distance)  # (layers, heads)

fig, ax = plt.subplots(figsize=(7, 3.6))
for layer in range(distance.shape[0]):
    ax.scatter(
        [layer] * distance.shape[1], distance[layer], s=18, color=COLORS["att"], alpha=0.8
    )
ax.set_xlabel("layer")
ax.set_ylabel("mean attention distance, pixels")
ax.set_title("Every dot is one head, averaged over the five photos")
plt.show()

Every dot is one head, averaged over the five photos

In the first layers the heads split: some look mostly at their neighbourhood, a patch or two away, like a convolution; others already look across the whole image, something a CNN can only do after many layers.

By the last layers every head looks far. The ViT paper reported the same spread; it is the clearest sign that the model learned a locality that a CNN has built in, and learned to drop it where it doesn’t help.

What decides the answer. The class token’s attention to the patches shows what it gathered, but a single layer tells little, because the token mixes with everything at every layer.

Attention rollout (Abnar and Zuidema, 2020) multiplies the attention matrices of all layers, each averaged over heads and mixed half and half with the identity for the residual connection, and reads the class token’s row of the product.

n = batch.shape[0]
rollout = torch.eye(grid * grid + 1).repeat(n, 1, 1)
for block in vit.blocks:
    a = block.attn.weights.mean(dim=1)  # (photos, 197, 197), averaged over heads
    a = 0.5 * a + 0.5 * torch.eye(grid * grid + 1)  # the residual path
    a = a / a.sum(-1, keepdim=True)
    rollout = a @ rollout
    block.attn.keep = False
maps = rollout[:, 0, 1:].reshape(n, grid, grid)

violet = to_rgb(COLORS["att"])
# from clear to violet, laid over the photo
glow = LinearSegmentedColormap.from_list("glow", [(*violet, 0.0), (*violet, 0.9)])
fig, axes = plt.subplots(1, n, figsize=(13, 2.9))
for k, photo in enumerate(photos.values()):
    heat = (maps[k] / maps[k].max()).numpy()
    axes[k].imshow(photo.resize((SIZE, SIZE)).convert("L"), cmap=greys, alpha=0.75)
    axes[k].imshow(
        heat, cmap=glow, extent=(0, SIZE, SIZE, 0), interpolation="bilinear"
    )
    axes[k].set_title(labels[ours[k].argmax().item()], fontsize=10)
    axes[k].axis("off")
plt.show()

Output of the code above

The violet regions are the bear’s head and the fish, the shuttle’s nose rather than its plume, the orca’s body, the hummingbird. The model was never told where the objects are; its answer simply rests on them.

Rollout is a rough tool, a product of averages that ignores the MLPs, but it is cheap and often right about where to look.

5. A bigger image, and stretched positions

More pixels per object help, and the ViT paper fine-tuned its models at 384×384. With the same 16-pixel patches the grid becomes 24 × 24 and the sequence 577 tokens long.

The patch embedding and the encoder don’t care, they work on any number of tokens; only the position embeddings, learned for a 14 × 14 grid, have to change.

The paper’s answer is lesson 9’s answer in two dimensions: interpolate them, treating the 196 vectors as a 14 × 14 image with 768 channels and resizing it.

def stretch_positions(pos, new_grid):
    cls_pos, grid_pos = pos[:, :1], pos[:, 1:]
    old_grid = int(grid_pos.shape[1] ** 0.5)
    d = grid_pos.shape[-1]
    grid_pos = grid_pos.reshape(1, old_grid, old_grid, d).permute(
        0, 3, 1, 2
    )  # (1, d, 14, 14)
    grid_pos = F.interpolate(
        grid_pos, size=(new_grid, new_grid), mode="bicubic", align_corners=False
    )
    grid_pos = grid_pos.permute(0, 2, 3, 1).reshape(1, new_grid * new_grid, d)
    return torch.cat([cls_pos, grid_pos], dim=1)


BIG = 384
big = ViT(vit_config, len(labels), image_size=BIG)
weights = rename(hf_weights)
weights["pos"] = stretch_positions(weights["pos"], BIG // PATCH)
big.load_state_dict(weights, strict=True)
big.eval()
big_batch = torch.stack([to_tensor(photo, size=BIG) for photo in photos.values()])
with torch.no_grad():
    big_probs = big(big_batch).softmax(-1)
print("tokens at 224:", grid * grid + 1, " at 384:", (BIG // PATCH) ** 2 + 1)
for k, name in enumerate(photos):
    small_answer = f"{labels[probs[k].argmax().item()]} {probs[k].max():.0%}"
    big_answer = f"{labels[big_probs[k].argmax().item()]} {big_probs[k].max():.0%}"
    print(f"{name:12} 224: {small_answer:28} 384: {big_answer}")

tokens at 224: 197  at 384: 577
bear         224: brown bear 100%              384: brown bear 100%
shuttle      224: space shuttle 100%           384: space shuttle 100%
lava         224: volcano 100%                 384: volcano 100%
orca         224: killer whale 100%            384: killer whale 100%
hummingbird  224: hummingbird 100%             384: hummingbird 100%

Without a single step of training, the model reads three times as many patches and still names every photo, as sure of itself as before.

These photos are easy; on harder images the stretched positions, close to but not the same as anything the model saw, are worth a short fine-tuning at the new size, which is what the ViT paper did.

The price of resolution is the length of the sequence, and attention’s cost grows with its square. The patch size sets the trade-off:

cost = []
d, layers = vit_config["hidden_size"], vit_config["num_hidden_layers"]
for size in [224, 384, 512]:
    for p in [32, 16, 14, 8]:
        n_tokens = (size // p) ** 2 + 1
        linear = (
            layers * n_tokens * 12 * d * d * 2
        )  # projections and the MLP, two flops a weight
        attention = layers * 2 * n_tokens * n_tokens * d * 2  # scores and the weighted sum
        cost.append(
            {
                "image": size,
                "patch": p,
                "tokens": n_tokens,
                "GFLOPs": (linear + attention) / 1e9,
                "attention's share": attention / (linear + attention),
            }
        )
table(pd.DataFrame(cost), digits=2)
imagepatchtokensGFLOPsattention's share
22432508.590.01
2241619734.890.04
2241425746.090.05
2248785156.060.15
3843214525.410.03
38416577110.290.11
38414730143.650.14
38482,305587.410.33
5123225746.090.05
512161,025212.850.18
512141,297282.330.22
51284,0971,314.730.47

At 224 pixels with 16-pixel patches attention is a small part of the work; at 512 pixels with 8-pixel patches it approaches half of it, and the share keeps growing with every step up.

This is why the vision encoders of today’s chat models cut images into tiles and why their patch sizes are 14 or 16 and not smaller: every halving of the patch quadruples the sequence.

6. Attention that finds objects without labels

Everything so far was learned from labels, a thousand ImageNet classes. In 2021 Mathilde Caron and colleagues at Facebook AI trained ViTs with no labels at all, by self-distillation: a student network learns to give the same output for a small crop of a photo that a slowly updated copy of itself, the teacher, gives for a large crop.

They called the method DINO, and noticed something nobody had asked for: the attention of the class token in the last layer outlined the objects in a photo, a segmentation learned from no segmentation at all.

DINO’s small ViT with 8-pixel patches is the same architecture as ours with other numbers, so it loads into our class. It was trained with the usual ImageNet normalisation, a different mean and deviation per channel.

DINO = "facebook/dino-vits8"
IMAGENET_MEAN, IMAGENET_STD = [0.485, 0.456, 0.406], [0.229, 0.224, 0.225]
if SMOKE:
    dino_config = {**vit_config, "patch_size": 8, "num_attention_heads": 4}
    dino = ViT(dino_config, 1)
else:
    dino_config = json.load(open(hf_hub_download(DINO, "config.json")))
    dino = ViT(dino_config, 1)
    dino_state = rename(load_file(hf_hub_download(DINO, "model.safetensors")))
    missing = dino.load_state_dict(dino_state, strict=False)
    print("only the classifier is missing:", missing.missing_keys)
dino.eval()

dino_batch = []
for photo in photos.values():
    dino_batch.append(to_tensor(photo, mean=IMAGENET_MEAN, std=IMAGENET_STD))
dino_batch = torch.stack(dino_batch)
last = dino.blocks[-1].attn
last.keep = True
with torch.no_grad():
    dino.encode(dino.tokens(dino_batch))
last.keep = False
side = SIZE // dino_config["patch_size"]
heads = last.weights[:, :, 0, 1:].reshape(
    len(photos), -1, side, side
)  # (photos, heads, 28, 28)
print("one map per head:", tuple(heads.shape))

only the classifier is missing: ['head.weight', 'head.bias']
one map per head: (5, 6, 28, 28)
shades = LinearSegmentedColormap.from_list("shades", [COLORS["card"], COLORS["att"]])
show = ["bear", "shuttle", "orca", "hummingbird"]
fig, axes = plt.subplots(len(show), heads.shape[1] + 1, figsize=(13, 2.1 * len(show)))
for r, name in enumerate(show):
    k = list(photos).index(name)
    axes[r, 0].imshow(photos[name].resize((SIZE, SIZE)))
    for h in range(heads.shape[1]):
        axes[r, h + 1].imshow(heads[k, h], cmap=shades, interpolation="nearest")
        if r == 0:
            axes[r, h + 1].set_title(f"head {h}", fontsize=10)
    for ax in axes[r]:
        ax.axis("off")
plt.show()

Output of the code above

Most heads of the last layer have found an object and outlined it on a grid of 8-pixel cells: the orca’s body and fins against the sea, the shuttle and the edge of its plume, the bear’s head and the fish.

Different heads take different parts, and one or two look at the background instead. The model was never shown a label, a box or a mask; separating a thing from its background turned out to be useful for matching crops of the same photo, and the attention learned to do it.

DINOv2, which Meta trained on 142 million curated images in 2023, is the same idea at scale, and its features are what many vision systems now start from. The last section puts them to work.

7. Land use from orbit, with few labels

EuroSAT is 27,000 images from the Sentinel-2 satellites of the European Space Agency’s Copernicus programme, 64×64 pixels, each 640 metres across, in ten classes of land use: crops, forest, pasture, highway, industry, housing, rivers, lakes and two more.

Helber and colleagues released it in 2019 under the MIT license. It is a good test of the ViT’s weak spot: images that look nothing like ImageNet, and a budget of labels we can shrink at will.

EUROSAT = "hf://datasets/timm/eurosat-rgb/data/{}-00000-of-00001.parquet"


def load_split(split, limit=None):
    frame = pd.read_parquet(EUROSAT.format(split))
    if limit:
        frame = frame.groupby("label").head(limit)
    images = []
    for record in frame["image"]:
        images.append(np.asarray(Image.open(io.BytesIO(record["bytes"])).convert("RGB")))
    images = torch.tensor(np.stack(images)).permute(0, 3, 1, 2)  # (n, 3, 64, 64), uint8
    names = [i.split("_")[0] for i in frame["image_id"]]
    return images, torch.tensor(frame["label"].values), names


train_images, train_labels, train_names = load_split("train", 30 if SMOKE else None)
test_images, test_labels, _ = load_split("test", 10 if SMOKE else None)
classes = {}
for label, name in zip(train_labels.tolist(), train_names):
    classes[label] = name
classes = [classes[k] for k in sorted(classes)]
print("train:", tuple(train_images.shape), " test:", tuple(test_images.shape))
print("classes:", ", ".join(classes))

fig, axes = plt.subplots(2, 10, figsize=(13, 3))
for c in range(10):
    picks = (train_labels == c).nonzero()[:2, 0]
    for r, i in enumerate(picks.tolist()):
        axes[r, c].imshow(train_images[i].permute(1, 2, 0))
        axes[r, c].axis("off")
    axes[0, c].set_title(classes[c], fontsize=8)
plt.show()

train: (16200, 3, 64, 64)  test: (5400, 3, 64, 64)
classes: AnnualCrop, Forest, HerbaceousVegetation, Highway, Industrial, Pasture, PermanentCrop, Residential, River, SeaLake

Output of the code above

Four learners, each trained on 10, 100 and all 1,620 labelled images per class:

  1. ViT-B/16 pretrained on ImageNet-21k’s 14 million images, the checkpoint before it learned ImageNet’s 1,000 classes, fine-tuned whole. The 64-pixel tiles are enlarged to 224.
  2. A small ViT from scratch: six blocks of width 192, 8-pixel patches on the 64-pixel tiles, so 64 patches.
  3. A small CNN from scratch with about as many parameters, four blocks of convolutions: the inductive bias the ViT gave up.
  4. A linear probe on DINOv2: the small DINOv2, frozen, turns each tile into a vector, and only a linear classifier on top is trained. No pixel-level training at all.

All four see the same augmentation: random flips and quarter turns, which change nothing about what a satellite tile shows.

def augment(images):
    flip = torch.rand(len(images), device=images.device) < 0.5
    images = torch.where(flip[:, None, None, None], images.flip(-1), images)
    turns = random.randint(0, 3)
    return torch.rot90(images, turns, dims=(-2, -1))


def normalise(images, size, mean, std):
    x = images.float() / 255
    if size != x.shape[-1]:
        x = F.interpolate(x, size=(size, size), mode="bilinear", align_corners=False)
    mean = torch.tensor(mean, device=x.device).reshape(1, -1, 1, 1)
    std = torch.tensor(std, device=x.device).reshape(1, -1, 1, 1)
    return (x - mean) / std


class SmallCNN(nn.Module):
    def __init__(self, classes, width=64):
        super().__init__()
        layers = []
        channels = 3
        for k in range(4):
            out = width * 2 ** min(k, 2)
            layers += [
                nn.Conv2d(channels, out, 3, padding=1),
                nn.BatchNorm2d(out),
                nn.GELU(),
            ]
            layers += [nn.Conv2d(out, out, 3, padding=1), nn.BatchNorm2d(out), nn.GELU()]
            layers += [nn.MaxPool2d(2)]
            channels = out
        self.features = nn.Sequential(*layers)
        self.head = nn.Linear(channels, classes)

    def forward(self, x):
        return self.head(self.features(x).mean(dim=(-2, -1)))


SMALL_VIT = {
    "hidden_size": 192,
    "num_hidden_layers": 6,
    "num_attention_heads": 3,
    "intermediate_size": 768,
    "patch_size": 8,
    "image_size": 64,
}
for name, net in [("small ViT", ViT(SMALL_VIT, 10)), ("small CNN", SmallCNN(10))]:
    print(f"{name}: {sum(p.numel() for p in net.parameters()) / 1e6:.1f}M parameters")

small ViT: 2.7M parameters
small CNN: 2.3M parameters

One training loop serves the first three learners. It counts steps rather than epochs, so that the smallest dataset is seen many times and the largest a few, and uses 16-bit floats with a loss scaler on the GPU.

def train(model, images, labels, size, mean, std, steps, lr, batch_size=64):
    model = model.to(device).train()
    optimizer = torch.optim.AdamW(model.parameters(), lr=lr, weight_decay=0.05)
    schedule = torch.optim.lr_scheduler.OneCycleLR(
        optimizer, lr, total_steps=steps, pct_start=0.1
    )
    scaler = torch.amp.GradScaler(device, enabled=device == "cuda")
    images, labels = images.to(device), labels.to(device)
    for _ in range(steps):
        pick = torch.randint(0, len(images), (batch_size,), device=device)
        x = normalise(augment(images[pick]), size, mean, std)
        with torch.autocast(device, dtype=torch.float16, enabled=device == "cuda"):
            loss = F.cross_entropy(model(x), labels[pick])
        optimizer.zero_grad()
        scaler.scale(loss).backward()
        scaler.step(optimizer)
        scaler.update()
        schedule.step()
    return model


@torch.no_grad()
def predict_all(model, images, size, mean, std, batch_size=256):
    model.eval()
    out = []
    for a in range(0, len(images), batch_size):
        x = normalise(images[a : a + batch_size].to(device), size, mean, std)
        with torch.autocast(device, dtype=torch.float16, enabled=device == "cuda"):
            out.append(model(x).float().argmax(-1).cpu())
    return torch.cat(out)


def subset(per_class):
    keep = []
    for c in range(10):
        keep += (train_labels == c).nonzero()[:per_class, 0].tolist()
    return train_images[keep], train_labels[keep]

DINOv2 comes from the transformers library: it adds a few details to the ViT recipe, such as layer scale and 14-pixel patches, and here it only serves as a frozen feature extractor. Its class token plus the average of its patch tokens is the image’s vector, and a linear layer on top is trained with plain gradient descent.

from transformers import AutoModel, Dinov2Config, Dinov2Model

if SMOKE:
    dinov2 = Dinov2Model(
        Dinov2Config(
            hidden_size=32,
            num_hidden_layers=1,
            num_attention_heads=2,
            intermediate_size=64,
            patch_size=14,
            image_size=56,
        )
    )
    dinov2_size = 56
else:
    dinov2 = AutoModel.from_pretrained("facebook/dinov2-small")
    dinov2_size = 224
dinov2 = dinov2.to(device).eval()


@torch.no_grad()
def dinov2_features(images, batch_size=256):
    out = []
    for a in range(0, len(images), batch_size):
        x = normalise(
            images[a : a + batch_size].to(device), dinov2_size, IMAGENET_MEAN, IMAGENET_STD
        )
        with torch.autocast(device, dtype=torch.float16, enabled=device == "cuda"):
            hidden = dinov2(pixel_values=x).last_hidden_state.float()
        out.append(torch.cat([hidden[:, 0], hidden[:, 1:].mean(1)], dim=-1).cpu())
    return torch.cat(out)


train_features = dinov2_features(train_images)
test_features = dinov2_features(test_images)
print("one tile as a DINOv2 vector:", tuple(train_features.shape[1:]))


def linear_probe(features, labels, steps=500):
    mean, std = features.mean(0), features.std(0) + 1e-6
    probe = nn.Linear(features.shape[1], 10)
    optimizer = torch.optim.AdamW(probe.parameters(), lr=1e-2, weight_decay=1e-3)
    x = (features - mean) / std
    for _ in range(steps):
        loss = F.cross_entropy(probe(x), labels)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
    with torch.no_grad():
        return probe((test_features - mean) / std).argmax(-1)

one tile as a DINOv2 vector: (768,)

The race. Every learner gets the same subsets of labelled tiles and is scored on the same 5,400 test tiles.

IN21K = "google/vit-base-patch16-224-in21k"
budgets = [3, 10] if SMOKE else [10, 100, 1620]
results = []
predictions = {}
for per_class in budgets:
    images, labels_ = subset(per_class)
    start = time.time()
    chosen = torch.cat([(train_labels == c).nonzero()[:per_class, 0] for c in range(10)])
    guess = linear_probe(train_features[chosen], train_labels[chosen])
    results.append(
        {
            "labels per class": per_class,
            "learner": "DINOv2 + linear probe",
            "accuracy": (guess == test_labels).float().mean().item(),
            "seconds": time.time() - start,
        }
    )
    predictions[("DINOv2 + linear probe", per_class)] = guess

    learners = [
        ("small CNN", lambda: SmallCNN(10), 64, IMAGENET_MEAN, IMAGENET_STD, 3e-3),
        ("small ViT", lambda: ViT(SMALL_VIT, 10), 64, IMAGENET_MEAN, IMAGENET_STD, 1e-3),
        ("ViT-B/16 pretrained", None, 224, [MEAN] * 3, [STD] * 3, 5e-5),
    ]
    for name, make, size, mean, std, lr in learners:
        start = time.time()
        if name == "ViT-B/16 pretrained":
            model = ViT(vit_config, 10)
            if not SMOKE:
                pretrained = rename(load_file(hf_hub_download(IN21K, "model.safetensors")))
                model.load_state_dict(pretrained, strict=False)  # a new head for 10 classes
            steps = 30 if SMOKE else max(300, 3 * len(images) // 32)
            model = train(model, images, labels_, size, mean, std, steps, lr, batch_size=32)
        else:
            steps = 30 if SMOKE else max(1500, 30 * len(images) // 64)
            model = train(make(), images, labels_, size, mean, std, steps, lr)
        guess = predict_all(model, test_images, size, mean, std)
        results.append(
            {
                "labels per class": per_class,
                "learner": name,
                "accuracy": (guess == test_labels).float().mean().item(),
                "seconds": time.time() - start,
            }
        )
        predictions[(name, per_class)] = guess
        del model
        if device == "cuda":
            torch.cuda.empty_cache()
results = pd.DataFrame(results)
order = ["ViT-B/16 pretrained", "small ViT", "small CNN", "DINOv2 + linear probe"]
accuracy = results.pivot(index="learner", columns="labels per class", values="accuracy")
accuracy = accuracy.loc[order]
accuracy.columns = [f"{n:,} labels per class" for n in accuracy.columns]
table(accuracy.reset_index(), digits=3)
learner10 labels per class100 labels per class1,620 labels per class
ViT-B/16 pretrained0.9160.9700.988
small ViT0.5860.8150.932
small CNN0.7400.9200.979
DINOv2 + linear probe0.8450.9260.966
fig, ax = plt.subplots(figsize=(7.2, 3.8))
for name, color in [
    ("ViT-B/16 pretrained", COLORS["att"]),
    ("DINOv2 + linear probe", COLORS["enc"]),
    ("small CNN", COLORS["data"]),
    ("small ViT", COLORS["dec"]),
]:
    part = results[results.learner == name]
    ax.plot(
        part["labels per class"] * 10,
        part.accuracy,
        marker="o",
        color=color,
        label=name,
        linewidth=2,
    )
ax.set_xscale("log")
ax.set_xlabel("labelled training tiles")
ax.set_ylabel("test accuracy")
ax.set_title("EuroSAT, ten classes: accuracy against the number of labels")
ax.legend(frameon=False, fontsize=9)
plt.show()

EuroSAT, ten classes: accuracy against the number of labels

The two from-scratch lines tell the story of 2020. With few labels the CNN is far ahead of the ViT of the same size: its built-in belief in local, repeated patterns is worth thousands of examples.

As labels grow, the ViT closes much of the gap, but on 16 thousand tiles it doesn’t pass the CNN; that took hundreds of millions of images.

The two pretrained learners show the way out the field took: do the data-hungry learning once, on a huge collection, and reuse it. A ViT that has seen ImageNet-21k, or a frozen DINOv2 with only a linear layer trained on top, is good with ten labels per class.

With all the labels, though, the CNN trained from scratch passes the frozen DINOv2: features that were never adapted to satellite colours have a ceiling, and the fine-tuned ViT, which adapts everything, stays on top.

Where the best model still goes wrong:

best = predictions[("ViT-B/16 pretrained", budgets[-1])]
confusion = torch.zeros(10, 10)
for t, p in zip(test_labels.tolist(), best.tolist()):
    confusion[t, p] += 1
confusion = confusion / confusion.sum(1, keepdim=True)
fig, ax = plt.subplots(figsize=(6.2, 5.4))
ax.imshow(confusion, cmap=shades, vmin=0, vmax=1)
ax.set_xticks(range(10), classes, rotation=60, fontsize=8)
ax.set_yticks(range(10), classes, fontsize=8)
ax.set_xlabel("predicted")
ax.set_ylabel("true")
for i in range(10):
    for j in range(10):
        if confusion[i, j] > 0.01:
            ink = "white" if confusion[i, j] > 0.5 else "black"
            ax.text(
                j,
                i,
                f"{confusion[i, j]:.2f}",
                ha="center",
                va="center",
                fontsize=7,
                color=ink,
            )
plt.show()

short = {  # titles that fit under a tile
    "AnnualCrop": "annual crop",
    "HerbaceousVegetation": "herbaceous",
    "PermanentCrop": "permanent crop",
    "SeaLake": "sea or lake",
}
wrong = (best != test_labels).nonzero()[:, 0][:12]
fig, axes = plt.subplots(2, 6, figsize=(12, 4.6))
for ax, i in zip(axes.flat, wrong.tolist()):
    ax.imshow(test_images[i].permute(1, 2, 0))
    truth = short.get(classes[test_labels[i]], classes[test_labels[i]].lower())
    guess = short.get(classes[best[i]], classes[best[i]].lower())
    ax.set_title(f"{truth}\n→ {guess}", fontsize=8)
    ax.axis("off")
for ax in axes.flat[len(wrong) :]:
    ax.axis("off")
plt.show()

Output of the code above

Output of the code above

The mistakes are few, and many of them involve herbaceous vegetation, taken for pasture, for crops or for forest: scrub, meadow and young trees are the same green-brown texture at this scale.

Annual and permanent crops are taken for each other too, and now and then a tile with a road across it is called a highway. Look at the tiles and you will hesitate too; at 10 metres a pixel, an orchard and a field of wheat can look the same.

8. Where the ViT is today

DeiT (Touvron and colleagues, Facebook AI, 2020) showed that with strong augmentation and regularisation a ViT trains on ImageNet alone, in under three days on one machine, and matches CNNs; the dependence on JFT was a matter of recipe as much as of architecture.

Self-supervised ViTs, DINO and DINOv2, learn from images without labels, and their features are what vision systems start from when labels are few, as section 7 showed.

Contrastive ViTs, CLIP and SigLIP, learn from images paired with captions; they are lesson 12.

The eyes of language models. A vision-language model takes a ViT’s output vectors, one per patch, projects them into the language model’s embedding space and puts them in the prompt next to the text tokens: lesson 13.

ViT-B/16’s passport, and the family tree with this lesson’s branch:

ViT-B/16's passport. Published October 2020 by Google Research, Brain Team, weights under Apache 2.0. Type: encoder-only Transformer for images, 86 million parameters. Input: 224 by 224 pixels cut into 196 patches of 16 by 16, plus a class token. Encoder: 12 blocks of width 768 with 12 heads and an MLP of 3,072, layer norm before each part. Positions: 197 learned vectors. Training: ImageNet-21k, 14 million images, then ImageNet-1k. Good at: classifying photos and serving as a backbone for other tasks. Not for: small datasets trained from scratch ViT-B/16's passport. Published October 2020 by Google Research, Brain Team, weights under Apache 2.0. Type: encoder-only Transformer for images, 86 million parameters. Input: 224 by 224 pixels cut into 196 patches of 16 by 16, plus a class token. Encoder: 12 blocks of width 768 with 12 heads and an MLP of 3,072, layer norm before each part. Positions: 197 learned vectors. Training: ImageNet-21k, 14 million images, then ImageNet-1k. Good at: classifying photos and serving as a backbone for other tasks. Not for: small datasets trained from scratch

The family of the Transformer as this course covers it, with lesson 11, the Vision Transformer, highlighted and lessons 1 to 10 done The family of the Transformer as this course covers it, with lesson 11, the Vision Transformer, highlighted and lessons 1 to 10 done

References

  1. Dosovitskiy, A., Beyer, L., Kolesnikov, A., et al. (2021). An image is worth 16x16 words: Transformers for image recognition at scale. ICLR 2021.
  2. Touvron, H., Cord, M., Douze, M., et al. (2021). Training data-efficient image transformers & distillation through attention. ICML 2021.
  3. Caron, M., Touvron, H., Misra, I., et al. (2021). Emerging properties in self-supervised vision transformers. ICCV 2021.
  4. Oquab, M., Darcet, T., Moutakanni, T., et al. (2024). DINOv2: learning robust visual features without supervision. Transactions on Machine Learning Research.
  5. Abnar, S., & Zuidema, W. (2020). Quantifying attention flow in transformers. ACL 2020.
  6. Naseer, M., Ranasinghe, K., Khan, S., et al. (2021). Intriguing properties of vision transformers. NeurIPS 2021.
  7. Helber, P., Bischke, B., Dengel, A., & Borth, D. (2019). EuroSAT: a novel dataset and deep learning benchmark for land use and land cover classification. IEEE Journal of Selected Topics in Applied Earth Observations and Remote Sensing.
  8. Krizhevsky, A., Sutskever, I., & Hinton, G. E. (2012). ImageNet classification with deep convolutional neural networks. NeurIPS 2012.
  9. Photographs, all public domain: Peter Pearsall, U.S. Fish and Wildlife Service, brown bear at Brooks Falls; NASA, Space Shuttle Discovery lifts off; J. D. Griggs, USGS, pāhoehoe fountain; NOAA Olympic Coast National Marine Sanctuary, orca breaching; U.S. Fish and Wildlife Service, rufous hummingbird at Seedskadee.
]]> Eva Koroleva How accurate is Polymarket? Calibration, longshot bias and a volume trap 2026-11-05T00:00:00+00:00 2026-11-05T00:00:00+00:00 https://blog.korolevalab.com/blog/2026/polymarket-accuracy

In short

Polymarket’s prices are good probabilities where people actually trade. A week before the close, markets with real trading are off by under two percentage points on average, and a day before the close they remove 82% of the error of a forecaster who knows only the base rate.

The famous accuracy figures say less than they seem, because most questions resolve No. What mispricing is left sits at the tails, cheap longshots a little too dear and near-certain favourites a little too cheap, and it is about the size of the spread.

1.8points between price and outcome, on average, a week before the close
77%right for a forecaster who answers No to every question
−39%average return on contracts bought at 2 to 5 cents

On the last evening of October I opened Polymarket to look at one number. The US midterm elections were three days away, and the market on which party would control the Senate had taken $13.5 million of bets.

A share that pays a dollar if Democrats win the Senate cost 71.5 cents. The House was hardly a contest by the market’s account: 94.5 cents for Democrats.

A week earlier the same kind of number had been wrong in public. Going into Brazil’s presidential runoff on 25 October, Polymarket and Kalshi gave Flávio Bolsonaro about an 80% chance of beating Lula, and Lula won[1]. Semafor’s headline that week was that prediction markets had “blown three big elections”[2].

By the time you read this you probably know how the Senate went, and I don’t. Suppose Republicans keep it. Was 71.5 wrong?

Suppose they lose it. Was it right? Neither follows, and Brazil does not prove the 80% wrong either.

A forecast of 71.5% claims something only about many situations like this one: that roughly seven in ten of them end the way the market leaned, and three in ten don’t.

So the question in the title cannot be settled by any one market, however famous. It can be settled by many. I downloaded the record of every market Polymarket resolved from 2024 on, almost three million of them, took the price histories of about 150,000 and checked what the prices said against what then happened.

This post is about what came out: how good the prices are as probabilities, where they are systematically off, a trap in the most natural way to slice the data, what the “accuracy” figures in the press leave out, and what a trader can and cannot do with any of it.

What a Polymarket price is, and why it should be a probability

Polymarket sells shares in outcomes. A share of “Yes” pays one dollar if the event happens and nothing if it does not; a share of “No” pays a dollar in the opposite case.

Since April 2026 the dollar is pUSD, Polymarket’s own token backed by the USDC stablecoin; before that it was USDC itself[3]. One Yes and one No together always pay exactly a dollar, so their prices add up to about a dollar too.

Every market is binary. A question with many answers, such as who wins a prize, is a set of Yes/No markets, one per candidate.

Shares trade in an order book, as on a stock exchange: buyers post bids, sellers post asks, and a trade happens when they meet.

The price the site shows, the number the press calls “the odds”, is the midpoint between the best bid and the best ask, or the last trade if the gap between them is wider than ten cents[4].

A Polymarket share on the US Senate on 31 October 2026: Yes, Democrats control the Senate, at 71.5¢, and No at 28.5¢; together they make a dollar. The best bid is 71¢ and the best ask 72¢. A Yes share bought at 72¢ pays $$1 if Democrats win, a gain of 28¢, and nothing if they lose, a loss of 72¢. It pays off on average only if the chance of a Democratic Senate is above 72% A Polymarket share on the US Senate on 31 October 2026: Yes, Democrats control the Senate, at 71.5¢, and No at 28.5¢; together they make a dollar. The best bid is 71¢ and the best ask 72¢. A Yes share bought at 72¢ pays $$1 if Democrats win, a gain of 28¢, and nothing if they lose, a loss of 72¢. It pays off on average only if the chance of a Democratic Senate is above 72%

Why should that price be a probability? Because of a simple argument about money. If you believe Democrats win the Senate with probability \(q\), a Yes share bought at price \(p\) is worth \(q \cdot 1 + (1 - q) \cdot 0 = q\) to you on average.

You buy when \(q > p\) and sell when \(q < p\). The price stops moving where nobody with money left is willing to push it either way.

That argument has holes, and economists have spent twenty years mapping them. The price is a compromise between traders, not anyone’s belief.

With logarithmic utility it equals the traders’ average belief, and for other reasonable preferences it stays close to it between about 20 and 80 cents[5]; for risk-neutral traders with different beliefs, a price of 75 cents is compatible with an average belief anywhere from 56% to 94%[6].

Money tied up in a bet until December could have earned interest. Since March 2026 most categories charge the taker a fee[7]. And the dollar arrives only if the event is resolved the way you expect, by a process with its own failures.

Whether the price works as a probability despite all that is an empirical question, and the rest of this post answers it with data.

How to check a forecaster

Calibration

A forecaster who said 70% is not wrong when the event fails to happen, but a forecaster who says 70% a thousand times should see about seven hundred of those events happen. That property is called calibration.

Weather services have been checked this way for seventy years: on the days a calibrated service gives a 30% chance of rain, it rains on about 30% of them.

The picture that shows calibration is the reliability diagram. Take all the forecasts, group them by value (all forecasts between 60% and 70%, say), and for each group plot the average forecast against the share of events that happened. A calibrated forecaster’s points fall on the diagonal.

If you have ever trained a classifier and checked whether its predicted probabilities can be trusted, you have drawn this picture already; modern neural networks are famously overconfident on it[8], and how to bin the forecasts is a small science of its own[9].

Calibration is not enough. A forecaster who always announces the base rate, the share of all events that happen, is perfectly calibrated and perfectly useless. About 34% of Polymarket’s markets resolve Yes, so saying “34%” about every market is calibrated: of all the markets you said it about, 34% happened.

What a good forecaster adds is resolution, forecasts that move away from the base rate in the right direction, high for the events that happen and low for those that don’t.

The Brier score

The Brier score measures both at once. For \(N\) forecasts \(p_i\) of events with outcomes \(y_i\) (1 if it happened, 0 if not), it is the mean squared error

\[\mathrm{BS} = \frac{1}{N} \sum_{i=1}^{N} (p_i - y_i)^2 .\]

A forecast of 70% scores \(0.3^2 = 0.09\) if the event happens and \(0.7^2 = 0.49\) if it does not. A forecaster who always says 50% scores \(0.25\), and a clairvoyant scores zero. Glenn Brier’s original score summed the error over both outcomes and ran from 0 to 2, so it is twice this one[10].

The score splits into a part for calibration, called reliability, a part for resolution, and a part that belongs to the questions rather than the forecaster. The skill score says how much of the worst honest score, that of the base-rate forecaster, a forecaster removes.

The Brier score in three parts: Murphy's decomposition

Allan Murphy showed that the score splits into three parts[11]. Group the forecasts into bins \(k\), each holding a share \(n_k / N\) of them, with average forecast \(\bar p_k\) and observed frequency \(\bar y_k\), and let \(\bar y\) be the overall base rate:

\[\mathrm{BS} \approx \underbrace{\sum_k \frac{n_k}{N} (\bar p_k - \bar y_k)^2}_{\text{reliability}} \;-\; \underbrace{\sum_k \frac{n_k}{N} (\bar y_k - \bar y)^2}_{\text{resolution}} \;+\; \underbrace{\bar y (1 - \bar y)}_{\text{uncertainty}} .\]

Reliability is the calibration error, the squared distance of the points from the diagonal, and should be small. Resolution is how far the groups’ outcomes spread away from the base rate, and should be large. Uncertainty belongs to the questions, not the forecaster: it is the score of the base-rate forecaster. The skill score \(1 - \mathrm{BS} / \bar y(1 - \bar y)\) says what share of that score a forecaster removes.

Three forecasters on the same markets

The three charts below put this on real data. All three forecasters see the same 18,157 markets a week before they closed. The first is Polymarket’s prices.

The second is a bolder copy that doubles the prices’ log-odds, so that 70% becomes 84% and 30% becomes 16%. The third is a timid copy that halves them, so that 70% becomes 60%. The grey square is the base-rate forecaster.

Three reliability diagrams for the same markets a week before they closed. Polymarket's own prices lie close to the diagonal, Brier score 0.079. A bolder copy, with the log-odds doubled, bends away from it, Brier 0.084; a timid copy, with them halved, bends the other way, Brier 0.089. A forecaster who always says the base rate, 25%, is a single point on the diagonal, Brier 0.186 Three reliability diagrams for the same markets a week before they closed. Polymarket's own prices lie close to the diagonal, Brier score 0.079. A bolder copy, with the log-odds doubled, bends away from it, Brier 0.084; a timid copy, with them halved, bends the other way, Brier 0.089. A forecaster who always says the base rate, 25%, is a single point on the diagonal, Brier 0.186

The market’s own prices score 0.079. Both copies score worse, 0.084 for the bold one and 0.089 for the timid one, and they fail in opposite ways.

The bold copy’s points lie above the diagonal on the left and below it on the right: when it says 5%, the event happens more often than that, because it took a modest forecast and inflated its confidence.

The timid copy bends the other way. Neither copy knows anything the market didn’t, and that is the point of the exercise. Distorting a calibrated forecaster in either direction only makes its score worse.

The base-rate forecaster says 25% about everything, the share of these markets that resolved Yes, and scores 0.186, more than twice the market’s error.

Why one election proves nothing

The case people quote as proof is the 2024 US presidential election. On the eve of the vote Polymarket priced a Trump win at 59 cents, while 538’s final model gave Kamala Harris 50 chances in 100[12]. Trump won all seven swing states.

Polymarket's price for Donald Trump winning the 2024 US presidential election over 2024: near 50 cents in spring, up after the first debate, down to about even after Biden withdrew and Harris entered, rising again in October; 1.53 billion dollars traded Polymarket's price for Donald Trump winning the 2024 US presidential election over 2024: near 50 cents in spring, up after the first debate, down to about even after Biden withdrew and Harris entered, rising again in October; 1.53 billion dollars traded

The market’s best-known trader that year, a Frenchman the press called Théo, had bet more than $30 million on Trump across several accounts.

He told the Wall Street Journal he had “absolutely no political agenda”[13], and later turned out to have commissioned his own polls, asking people not whom they would vote for but whom their neighbours would[14]. He made tens of millions of dollars, and he was right.

Think about what that shows. A forecaster who put 59% on the winner scores \(0.41^2 \approx 0.17\) for that election; one who put 50% scores \(0.25\).

A single draw from a coin that lands heads 59% of the time says almost nothing about the coin, and the same arithmetic run on Brazil a week ago says almost nothing in the other direction. One outcome cannot separate a good forecaster from a lucky one.

Thousands can.

Three million markets, and which of them are forecasts

Polymarket publishes everything needed to check it. Its Gamma API lists every event and market with its question, its tags, the dollars traded and the final outcome; its order-book API returns the price history of every outcome token.

I took every event that closed from 2024 to the end of October 2026 with at least $1,000 traded: 2,959,534 binary markets that resolved cleanly to Yes or No. (The API records only a few thousand markets before 2024, and I leave them out of the counts.)

About 0.4% of markets with more than $10,000 traded resolved 50-50, the platform’s answer to a question that turned out ambiguous, and I set those aside.

What Polymarket has become

Stacked bars of resolved Polymarket markets by the year they closed and topic: 2024: 12,964, 2025: 193,734, 2026: 2,752,836. 34% of all of them resolved Yes Stacked bars of resolved Polymarket markets by the year they closed and topic: 2024: 12,964, 2025: 193,734, 2026: 2,752,836. 34% of all of them resolved Yes

Polymarket in 2026 is mostly a sports book. One football match can carry a dozen markets (the winner, the handicap, the total goals), and sports make up 73% of all the markets that closed this year.

Next come the short crypto markets that ask whether Bitcoin will be higher at the end of the next fifteen minutes, hour or day, 11%, and science and weather questions, 6%. Politics, the subject of nearly every article about prediction markets, is 2% of the count.

Downloading the price history of three million markets is not practical, so I sampled about 150,000 of them, and how to sample turned out to matter more than anything else in this study.

Big markets were drawn more often than small ones, and every pooled figure is weighted back to the whole population; a later section shows what goes wrong without that.

How the markets were drawn and weighted

I made two independent random draws. The first takes a share of the markets in each topic, whatever their size: 1.7% of sports, 11% of politics, 16% of economics, 20% of culture, 5.5% of crypto, 3.5% of science and weather.

The second takes every market with more than $100,000 traded (every sports market with more than a million) and a share of the smaller ones above $10,000. Together they hold 153,602 markets in 80,320 events.

Every market had a known chance of being drawn, and when I pool markets each counts in inverse proportion to that chance, so the pooled figures describe all resolved markets rather than my sample; statisticians call this the Horvitz–Thompson estimator[15].

Drawing big markets more often is harmless as long as the weights undo it. Leaving markets out, or grouping them by size, is not. Uncertainty comes from a bootstrap over events rather than markets, because the twenty candidates of one election are not twenty independent questions.

Market lists come from the Gamma API (events/keyset), prices from the order-book API (prices-history, an hourly series for the last two weeks of a market and a 12-hour one for its whole life), and trades from the data API (trades).

What counts as a forecast

The forecast I score is the price of a market’s first outcome (Yes, in a Yes/No market) a fixed time before the market closed: a week, mostly, and a day, a month or three months where I compare horizons.

The close comes after the outcome is known, typically hours after a match ends or a day or two after an election, so a week before the close is before the event itself.

A market counts as a forecast only if at least $1,000 had been traded in it in the week before. This matters more than it sounds. Most markets on Polymarket are created, barely traded and resolved, and the price history of such a market is the middle of an almost empty order book.

One football market in my sample, on whether a Bundesliga match would have five goals or more, stood at 48 cents a week before the game; most matches end with fewer, this one did too, and over its whole life the market traded $134. That price is not a forecast of a high-scoring game.

A bid at one cent and an ask at 95 cents would give it. The test uses only what was known at the time, the trades of that week, which I took from Polymarket’s public trades API.

Calibration: a week out, prices are off by 1.8 points

Here is the main result.

Reliability diagram: markets grouped by their price a week before closing, against the share that resolved Yes: priced 1%, happened 0%; priced 3%, happened 2%; priced 7%, happened 9%; priced 14%, happened 15%; priced 25%, happened 28%; priced 35%, happened 36%; priced 45%, happened 46%; priced 55%, happened 60%; priced 65%, happened 70%; priced 75%, happened 82%; priced 85%, happened 89%; priced 93%, happened 92%; priced 97%, happened 99%; priced 99%, happened 100%. Next to it, a histogram of those prices Reliability diagram: markets grouped by their price a week before closing, against the share that resolved Yes: priced 1%, happened 0%; priced 3%, happened 2%; priced 7%, happened 9%; priced 14%, happened 15%; priced 25%, happened 28%; priced 35%, happened 36%; priced 45%, happened 46%; priced 55%, happened 60%; priced 65%, happened 70%; priced 75%, happened 82%; priced 85%, happened 89%; priced 93%, happened 92%; priced 97%, happened 99%; priced 99%, happened 100%. Next to it, a histogram of those prices

The points sit close to the diagonal. Of the 18,157 markets with real trading a week before the close, those priced around 25 cents resolved Yes 28% of the time and those around 75 cents 82%.

The average gap between price and frequency is 1.8 percentage points, and the reliability term of the Brier score is 0.001, a hundredth of the resolution term, 0.108. A weather service would be pleased with that.

When a liquid market says 70%, read it as about 70%, and remember that 70 means it fails three times in ten.

The small deviations have a shape. Between 55 and 85 cents events happened a little more often than priced (60% of the markets at 55 cents, 82% of those at 75), while contracts at 2 to 5 cents happened less often than priced, 2.3% of the time against a price of 3.2. The prices are slightly compressed toward the middle: favourites a little too cheap, longshots a little too dear.

That is the pattern economic theory predicts for traders who like long shots or dislike tying up money, and it is what a Princeton thesis on 188,509 Polymarket markets found too[16].

The right panel shows something just as important: where the forecasts are. Almost half of these markets are priced below 5 cents a week before they close.

Most questions on Polymarket are of the form “will this specific thing happen”, most of them don’t happen, and the market knows. Those confident, correct and uninteresting forecasts are a large part of every headline accuracy figure.

Accuracy against a forecaster who always says no

Polymarket advertises its record on a page titled “How accurate is Polymarket?”: 98.7% accurate four hours before resolution, 90.3% a month before[17]. Accuracy there means the share of markets where the leading outcome won.

It sounds like a hard bar to clear. The same page notes that 78.3% of its markets resolve No, and explains why: “There are usually more ways for something not to happen than to happen in one exact way”.

A forecaster who answers “No” to every question without reading it is therefore 78% accurate on Polymarket’s own sample. The left half of the next chart shows what that does to the comparison on mine.

The same 5,420 markets a month, a week and a day before they closed. Left: share where the side priced above 50 cents won, against the share a forecaster who always says No gets right: 1 month before 87% against 77%; 1 week before 93% against 77%; 1 day before 96% against 77%. Right: Brier skill over the base rate: 1 month 48% for all, 40% for those priced 2 to 98 cents; 1 week 70% for all, 54% for those priced 2 to 98 cents; 1 day 82% for all, 57% for those priced 2 to 98 cents The same 5,420 markets a month, a week and a day before they closed. Left: share where the side priced above 50 cents won, against the share a forecaster who always says No gets right: 1 month before 87% against 77%; 1 week before 93% against 77%; 1 day before 96% against 77%. Right: Brier skill over the base rate: 1 month 48% for all, 40% for those priced 2 to 98 cents; 1 week 70% for all, 54% for those priced 2 to 98 cents; 1 day 82% for all, 57% for those priced 2 to 98 cents

The chart uses the 5,420 markets that had real trading a month, a week and a day before they closed, so the same questions are scored at every horizon.

The favourite won 87% of them a month before the close, 93% a week before and 96% a day before. A forecaster who says No to everything gets 77% right on these markets at every horizon, since only 23% of them resolved Yes.

The honest comparison is skill: the market removes 48% of the base-rate forecaster’s Brier score a month out, 70% a week out and 82% a day out. On the questions still in doubt at the time, priced between 2 and 98 cents, it removes 40%, 54% and 57%.

Key point

Accuracy flatters any forecaster whose questions mostly resolve one way. Ask for skill: the share of the base-rate forecaster’s error that is gone. On Polymarket it is 70% a week before the close, and 54% on the questions still in doubt.

Those are the numbers to keep in mind the next time anyone, a platform or a pundit, advertises 90% accuracy.

How the forecast sharpens as the close approaches

Information arrives over time. The Brier score of markets that lived at least two months traces how quickly.

Brier score of 3,492 markets, each open for at least 60 days, by days before the close: 0.084 at 60 days, 0.056 at 30, 0.033 at 7 and 0.021 at half a day, against 0.150 for always saying the base rate. The share priced below 1 cent or above 99 rises from 26% to 75% Brier score of 3,492 markets, each open for at least 60 days, by days before the close: 0.084 at 60 days, 0.056 at 30, 0.033 at 7 and 0.021 at half a day, against 0.150 for always saying the base rate. The share priced below 1 cent or above 99 rises from 26% to 75%

Sixty days before the close, the 3,492 markets that lived that long and were traded a month out already score 0.084, well below the base rate’s 0.150.

The score falls steadily, to 0.056 a month out and 0.033 a week out, and drops again in the last days, when scheduled events (a vote, a match, a data release) take place.

The lower panel counts the markets already priced at the extremes, below a cent or above 99 cents: a quarter of them at sixty days, more than half a week out, three quarters at the end. A market at 99.5 cents is not a forecast any more; it is a receipt waiting to be cashed.

The volume trap: why big markets seem to underprice longshots

The obvious next question is whether big markets forecast better than small ones, and the obvious way to answer it is to group markets by the dollars traded. Polymarket’s accuracy page does exactly this. It is a trap, and I walked into it first.

My first sample took every market with more than $100,000 traded and a share of those between $10,000 and $100,000, the way one naturally filters out dead markets.

In that sample, contracts priced between 10 and 20 cents a week before the close returned 25% on average, and the bigger the market, the more the longshots won.

It looked like the reverse of everything the racetrack literature says, and a trading strategy in one line. Here are the longshots of the sample drawn by topic, grouped two ways.

Contracts priced 2 to 20 cents a week before the close: how often they won divided by their price, grouped by the market's total volume and by the dollars traded in it in the week before. By total volume: under $$1K 0.49; $$1K–10K 0.60; $$10K–100K 1.13; $$100K–1M 1.34; $$1M–10M 2.14. By the week before: under $$100 0.72; $$100–1K 0.87; $$1K–10K 0.96; $$10K–100K 0.81 Contracts priced 2 to 20 cents a week before the close: how often they won divided by their price, grouped by the market's total volume and by the dollars traded in it in the week before. By total volume: under $$1K 0.49; $$1K–10K 0.60; $$10K–100K 1.13; $$100K–1M 1.34; $$1M–10M 2.14. By the week before: under $$100 0.72; $$100–1K 0.87; $$1K–10K 0.96; $$10K–100K 0.81

On the left, grouped by the dollars a market traded over its whole life, contracts priced 2 to 20 cents won about half as often as priced in markets under $1,000 (a ratio of 0.49) and twice as often in markets of $1–10 million (2.14).

On the right, the same contracts grouped by the dollars traded in the week before the forecast show nothing of the kind: 0.72 in the thinnest markets, 0.87 and 0.96 as trading grows. Favourites mirror it, from 1.05 to 0.88 on the left and flat at about 1.0 on the right.

The left panel is an artefact, and the reason is in what volume counts. A market’s total volume includes all the trading after the forecast, and nothing makes people trade like a surprise.

When a 10-cent candidate starts winning, the market fills with buyers, sellers, hedgers and spectators, and its volume jumps; when the favourite wins as expected, the market stays quiet and small.

a longshot starts winningeveryone trades itits volume grows"big markets underprice longshots"

So sorting markets by total volume quietly sorts them by whether the unlikely happened. In statistics this is selection on a consequence of the outcome, the same family of error as look-ahead bias in a backtest, which I wrote about in the post on AI crypto trading bots. Polymarket’s own accuracy page breaks its Brier scores down by volume in exactly this way.

Trap

Never filter, group or weight markets by their total volume, or by anything else that is counted after the forecast. Use what was known when the price was set: the dollars traded in the week before, the order book at the time, the topic.

The right panel is the honest version, because the week before a forecast is over when the forecast is made. It says something useful.

In thin markets longshots are overpriced, and the more a market trades, the closer to fair they become. That is why the rest of this post, like the calibration above, uses the test of $1,000 traded in the week before and never groups markets by their total volume.

Longshots lose, favourites earn a little

Calibration is an average, and averages hide the tails. The tails are where traders make or lose money, and economists have known for seventy years where to look: at the racetrack, bettors overpay for longshots and underpay for favourites[18].

In six million American horse races, betting on horses at 100 to 1 or longer lost about 61% of the stake, while betting on the favourite in every race lost only 5.5%[19]. On Kalshi, contracts bought below 10 cents lost over 60% of the money put into them[20].

To look for it on Polymarket, I treated both sides of every market as a contract: a Yes at the price \(p\) and a No at \(1 - p\). Buying a contract at price \(p\) that pays \(y \in \{0, 1\}\) returns \(y / p - 1\) on the stake. If prices were exact probabilities, the average return in every price band would be zero.

Average return of buying a contract a week before the close, by price band: 0.5–2¢ -22.5%; 2–5¢ -38.8%; 5–10¢ +19.4%; 10–20¢ -2.5%; 20–30¢ -4.1%; 30–40¢ -3.8%; 40–50¢ -3.9%; 50–60¢ +3.3%; 60–70¢ +2.9%; 70–80¢ +0.3%; 80–90¢ +0.0%; 90–95¢ -1.4%; 95–98¢ +1.1%; 98–99.5¢ +0.3% Average return of buying a contract a week before the close, by price band: 0.5–2¢ -22.5%; 2–5¢ -38.8%; 5–10¢ +19.4%; 10–20¢ -2.5%; 20–30¢ -4.1%; 30–40¢ -3.8%; 40–50¢ -3.9%; 50–60¢ +3.3%; 60–70¢ +2.9%; 70–80¢ +0.3%; 80–90¢ +0.0%; 90–95¢ -1.4%; 95–98¢ +1.1%; 98–99.5¢ +0.3%

The bias is there, and it lives at the extremes. Contracts priced between 2 and 5 cents a week before the close won 2.1% of the time while costing 3.2 cents on average, so buying them lost 39% of the stake, with a 95% interval from −53% to −22%.

The cheapest band, under 2 cents, lost 22%, with an interval wide enough to include zero.

At the other end, contracts at 95 to 98 cents won 97.8% of the time and returned 1.15% in the week, and those at 98 to 99.5 cents returned 0.35%; both intervals stay above zero. Between 5 and 95 cents the averages cannot be told apart from zero.

Longshots under 5 cents lose a quarter to two fifths of the stake; favourites above 95 cents earn a fraction of a percent a week, before costs.

How the returns are counted, and why the weighting matters

How contracts are weighted matters for this kind of result. A study of 588 million Polymarket trades found that buys under 10 cents lost 19.3 cents per dollar when weighted by the money traded, 6.3 cents when each contract counted once, and gained 4.1 cents when related contracts were first grouped by event[21].

My numbers count each market once, weighted to the population of markets, with uncertainty from resampling whole events.

A return here is that of one contract bought at the mid price a week before the close and held to the end, averaged over contracts in a price band with the population weights; the interval resamples whole events 400 times.

A favourite at 97 cents that earns a percent in a week sounds like free money, and it is not. The bars are mid prices, and nobody buys at the mid.

The downside is a total loss on the rare occasion the favourite fails, and the 50-50 resolutions I excluded would cost a favourite buyer half the stake; at about 0.4% of markets they take back up to 0.2 points of that return.

What the spread and the fees take

The price shown is the middle of the order book, but a buyer pays the ask. On the last night of October I recorded the best bid and ask of all 124,068 open markets, 57,468 of them with orders on both sides, and measured what buying at the ask costs above the middle, as a share of the price paid.

Median cost of buying at the ask instead of the middle of the book, as a share of the price, by contract price: 0.5–2¢ 66.7%; 2–5¢ 66.1%; 5–10¢ 29.4%; 10–20¢ 13.0%; 20–30¢ 6.7%; 30–40¢ 4.8%; 40–50¢ 2.5%; 50–60¢ 2.0%; 60–70¢ 2.4%; 70–80¢ 2.1%; 80–90¢ 2.3%; 90–95¢ 2.2%; 95–98¢ 2.1%; 98–99.5¢ 0.6% Median cost of buying at the ask instead of the middle of the book, as a share of the price, by contract price: 0.5–2¢ 66.7%; 2–5¢ 66.1%; 5–10¢ 29.4%; 10–20¢ 13.0%; 20–30¢ 6.7%; 30–40¢ 4.8%; 40–50¢ 2.5%; 50–60¢ 2.0%; 60–70¢ 2.4%; 70–80¢ 2.1%; 80–90¢ 2.3%; 90–95¢ 2.2%; 95–98¢ 2.1%; 98–99.5¢ 0.6%

For contracts between 40 and 95 cents the median cost is about 2% of the price. Below 20 cents it explodes: 13% for contracts at 10 to 20 cents, 29% at 5 to 10, and about two thirds of the price under 5 cents.

A two-cent contract with a bid at one cent and an ask at three costs a full cent above its middle, half of what it is worth on paper. At the other end, near-certain contracts above 98 cents are cheap to buy, 0.6% of the price.

On top of the spread comes the fee. Since 30 March 2026 a taker on most categories pays \(r \cdot p(1 - p)\) per share, where \(p\) is the price and the rate \(r\) is 0.04 for politics, 0.05 for sports and most other categories and 0.07 for crypto; makers pay nothing, and geopolitical markets are free[7].

As a share of the money paid that is \(r(1 - p)\): about 5% of the stake for a 5-cent sports contract, a quarter of a percent for a 95-cent one.

Many of the busiest markets ran with the fee switched off on the day I looked, so the rate has to be checked market by market.

Put together, a longshot buyer who crosses the spread gives up between a tenth and two thirds of the stake before the event even happens.

A favourite buyer gives up a couple of percent, more than the whole edge the previous chart found: after half the spread the 95–98 cent band goes from +1.15% to −0.9% a week, and the 98–99.5 cent band from +0.35% to −0.3%. The diamonds on that chart show it.

For traders

The longshot discount and the favourite premium are real, but they are about the size of what it costs to trade them, which is why they survive. Selling longshots and buying near-certainties works only for those who post limit orders and let others cross the spread.

Topics: politics is calibrated best

Brier skill and calibration error a week before the close, by topic: Politics and world skill 58%, error 1.9 points; Sports skill 53%, error 4.4 points; Crypto prices and events skill 54%, error 3.9 points; Economy and markets skill 60%, error 3.3 points; Culture and media skill 59%, error 2.3 points; Science, tech, weather skill 70%, error 4.1 points Brier skill and calibration error a week before the close, by topic: Politics and world skill 58%, error 1.9 points; Sports skill 53%, error 4.4 points; Crypto prices and events skill 54%, error 3.9 points; Economy and markets skill 60%, error 3.3 points; Culture and media skill 59%, error 2.3 points; Science, tech, weather skill 70%, error 4.1 points

A week before the close, markets with real trading are skilled in every topic, from about 53% of the base rate’s error removed in sports and crypto to 70% in science and weather.

Politics, with 8,929 markets the largest group, is also the best calibrated, with an average gap of 1.9 points; sports and science have the largest gaps, 4.4 and 4.1, with crypto close behind at 3.9. Sports here means the markets traded a week before the close, mostly season futures (who wins the league, who wins the award) rather than single games, which are listed only days before kick-off.

Crypto markets resolve together

Crypto deserves a chart of its own, because its errors have a cause that has nothing to do with the traders’ judgement.

Two panels by month. Top: among crypto price markets priced 30 to 70 cents a week before the close, the share where the side betting on rising prices won, against that side's average price. Bottom: Bitcoin's return that month. The two move together: 2024-12: 29% of 17, Bitcoin -3%; 2025-01: 58% of 12, Bitcoin +9%; 2025-02: 39% of 18, Bitcoin -18%; 2025-03: 31% of 16, Bitcoin -2%; 2025-04: 67% of 12, Bitcoin +14%; 2025-05: 73% of 15, Bitcoin +11%; 2025-06: 40% of 10, Bitcoin +2%; 2025-07: 75% of 16, Bitcoin +8%; 2025-09: 60% of 10, Bitcoin +5%; 2025-10: 20% of 10, Bitcoin -4%; 2025-11: 29% of 14, Bitcoin -18%; 2026-01: 41% of 17, Bitcoin -10%; 2026-02: 30% of 10, Bitcoin -15%; 2026-04: 67% of 15, Bitcoin +12%; 2026-09: 50% of 10, Bitcoin +6% Two panels by month. Top: among crypto price markets priced 30 to 70 cents a week before the close, the share where the side betting on rising prices won, against that side's average price. Bottom: Bitcoin's return that month. The two move together: 2024-12: 29% of 17, Bitcoin -3%; 2025-01: 58% of 12, Bitcoin +9%; 2025-02: 39% of 18, Bitcoin -18%; 2025-03: 31% of 16, Bitcoin -2%; 2025-04: 67% of 12, Bitcoin +14%; 2025-05: 73% of 15, Bitcoin +11%; 2025-06: 40% of 10, Bitcoin +2%; 2025-07: 75% of 16, Bitcoin +8%; 2025-09: 60% of 10, Bitcoin +5%; 2025-10: 20% of 10, Bitcoin -4%; 2025-11: 29% of 14, Bitcoin -18%; 2026-01: 41% of 17, Bitcoin -10%; 2026-02: 30% of 10, Bitcoin -15%; 2026-04: 67% of 15, Bitcoin +12%; 2026-09: 50% of 10, Bitcoin +6%

Most crypto markets ask about one number, the price of Bitcoin or Ether on a date. Take the close calls, markets priced between 30 and 70 cents a week out, and orient each so that Yes means the price went up.

In the months when Bitcoin rose, the bet on rising prices won 63% of the time; in the months when it fell, 32%. Its price a week before was about 50 cents in both. Across fifteen months, the share of winning bullish bets and Bitcoin’s monthly return have a correlation of 0.78.

The market was not wrong in any useful sense: a week before the close nobody knew which way Bitcoin would go, and 50 cents was a fair price for a coin toss. But a dozen such markets in one month are not a dozen tosses. They are one toss recorded a dozen times.

A calibration chart of crypto markets mostly shows whether the period studied was a good one for Bitcoin. Calibration needs independent questions, which is why I resample whole events rather than markets, and why even that understates the uncertainty for crypto.

Long-dated favourites and the price of waiting

A market that resolves in six months has to pay its traders for waiting. A No share at 95 cents on a question that will almost certainly resolve No pays 5 cents on 95 after half a year, about 10.8% a year, if it pays at all; three-month US Treasury bills paid 4.07% on average in 2025 and 3.71% so far in 2026[22].

Traders who could earn interest elsewhere will not bid near-certain long-dated contracts all the way to a dollar, and the people with the most money to commit, favourite buyers, drop out first.

Page and Clemen found the result on the old Intrade exchange: markets far from expiry were poorly calibrated, with prices “systematically biased toward 0.50”[23].

Polymarket noticed. Since September 2025 it has paid interest on positions in a short list of long-dated markets “to keep long-term pricing accurate”, 4% a year at launch and 3.25% since June 2026[24].

The list holds about eighteen events, the 2028 presidential race and the midterms among them. Every other market pays nothing for waiting.

Contracts priced above 90 cents one, three and six months before the close: their price, how often they won and what holding them to the end returned per year: 1 month 90–95¢: price 92.6%, won 94.3%, +25.0% a year; 1 month 95–98¢: price 96.6%, won 98.4%, +23.9% a year; 1 month 98–99.5¢: price 98.8%, won 99.3%, +6.4% a year; 3 months 90–95¢: price 92.6%, won 94.7%, +9.3% a year; 3 months 95–98¢: price 96.7%, won 98.1%, +6.2% a year; 3 months 98–99.5¢: price 98.9%, won 97.7%, -4.9% a year; 6 months 90–95¢: price 92.7%, won 97.0%, +9.3% a year; 6 months 95–98¢: price 96.6%, won 93.3%, -6.8% a year; 6 months 98–99.5¢: price 98.9%, won 98.7%, -0.4% a year Contracts priced above 90 cents one, three and six months before the close: their price, how often they won and what holding them to the end returned per year: 1 month 90–95¢: price 92.6%, won 94.3%, +25.0% a year; 1 month 95–98¢: price 96.6%, won 98.4%, +23.9% a year; 1 month 98–99.5¢: price 98.8%, won 99.3%, +6.4% a year; 3 months 90–95¢: price 92.6%, won 94.7%, +9.3% a year; 3 months 95–98¢: price 96.7%, won 98.1%, +6.2% a year; 3 months 98–99.5¢: price 98.9%, won 97.7%, -4.9% a year; 6 months 90–95¢: price 92.7%, won 97.0%, +9.3% a year; 6 months 95–98¢: price 96.6%, won 93.3%, -6.8% a year; 6 months 98–99.5¢: price 98.9%, won 98.7%, -0.4% a year

Polymarket shows the same thing. Contracts priced between 90 and 95 cents a month before the close won 94.3% of the time against a price of 92.6, and those at 95 to 98 cents won 98.4% against 96.6.

Held to the end, both earned about 1.8% in a month, which compounds to roughly 25% a year, six times what a Treasury bill paid.

Three months out, the same two bands earned 6 to 9% a year and the most expensive band lost money; six months out the samples are small and the results noisy.

A month before the close, favourites at 90 to 98 cents paid about 25% a year at mid prices: the price of waiting, not of ignorance.

The pattern fits Page and Clemen: near-certain long-dated contracts are a little too cheap, and the discount is about what impatient money charges. It is not free money either.

These are mid prices, a single failure costs the whole stake, and a trader who buys favourites a month out is lending to the market at a rate that has to cover the risk that the near-certain doesn’t happen.

A two-number correction that doesn’t help

An engineer looking at the calibration chart will ask the obvious question: if the market is slightly off in a systematic way, can a model correct it? This is the same move as calibrating a classifier. Fit two numbers, \(a\) and \(b\), so that a corrected probability

\[\hat q = \sigma\big(a \cdot \operatorname{logit}(p) + b\big), \qquad \operatorname{logit}(p) = \ln \frac{p}{1 - p}, \quad \sigma(z) = \frac{1}{1 + e^{-z}},\]

fits the outcomes better. With \(b = 0\), a slope \(a > 1\) means the prices should have been bolder, longshots cheaper and favourites dearer, and \(a < 1\) means they should have been more timid. This is Platt scaling, a standard way to recalibrate a classifier[25].

A Princeton thesis that fitted it to 188,509 Polymarket markets found a slope of 1.112[16], and a study of Kalshi and Polymarket found political markets “compressed toward 50%” in the same sense[26].

The interesting question is whether such a correction helps on markets it has not seen. I fitted \(a\) and \(b\) only on markets that closed before July 2025 and scored the correction on markets that closed after.

The fitted correction from Polymarket's price to a probability, against the diagonal, for one week and one month before the close: 1 week: slope 1.12, Brier on later markets 0.0796 before and 0.0797 after; 1 month: slope 1.18, Brier on later markets 0.0758 before and 0.0761 after The fitted correction from Polymarket's price to a probability, against the diagonal, for one week and one month before the close: 1 week: slope 1.12, Brier on later markets 0.0796 before and 0.0797 after; 1 month: slope 1.18, Brier on later markets 0.0758 before and 0.0761 after

Fitted on the markets that closed before July 2025, the correction has a slope of 1.12 a week before the close and 1.18 a month before: the prices should have been a little bolder, the same compression the calibration chart showed, and close to Walker’s 1.112.

Applied to the 13,618 markets that closed later, it changes the Brier score a week out by −0.0001, with a 95% interval from −0.0009 to +0.0008, and a month out by −0.0003. In other words, nothing.

The bias is real on average, but too small and too unstable over time for a formula to turn it into a better forecast, let alone a profit after the spread.

For engineers

The same discipline applies to any recalibration of a model: fit it on the past, score it on the future, and report the difference with its interval. Here a correction that every in-sample diagnostic recommends turns out to be worth nothing.

Insiders: when the market knows too early

The most impressive forecasts on Polymarket are sometimes not forecasts at all.

The Nobel Peace Prize

Prices of the leading candidates for the 2025 Nobel Peace Prize on Polymarket from midday on 8 October to the announcement on 10 October. María Corina Machado's price is near the bottom until late on 9 October, then rises within hours to about two thirds of a dollar before the announcement at 09:00 UTC Prices of the leading candidates for the 2025 Nobel Peace Prize on Polymarket from midday on 8 October to the announcement on 10 October. María Corina Machado's price is near the bottom until late on 9 October, then rises within hours to about two thirds of a dollar before the announcement at 09:00 UTC

The Norwegian Nobel Committee announced the 2025 Peace Prize at 09:00 UTC on 10 October. For months María Corina Machado had traded at a cent or two. Shortly after midnight in Oslo her price started to climb, passed 43 cents before one in the morning and reached 73 by two[27].

In the price history I downloaded she was still at 67 cents five minutes before the announcement. An account created days earlier bought early and made tens of thousands of dollars.

The director of the Norwegian Nobel Institute said, “We take this very seriously”[28]. In January the institute concluded that the winner’s name had most likely leaked through “the digital domain” and that it could not say who was behind it[29].

Maduro, the Iran strikes and Google

That was not the only case. On the night of 2 January 2026, hours before Donald Trump announced the capture of Nicolás Maduro, a weeks-old account bet about $33,000 that Maduro would be out by the end of the month and made about $400,000.

In April the Justice Department charged a Special Forces master sergeant who had helped plan the operation, and the CFTC called it the first insider-trading case it had ever brought over event contracts[30, 31]. He has pleaded not guilty.

Hours before the US strikes on Iran on 28 February, six freshly funded wallets bought Yes on the strike market and made about a million dollars[32].

In May a Google engineer was charged over about $1.2 million made on markets about Google’s “Year in Search” lists, using, according to Google, marketing material he could see at work[33].

Polymarket’s chief executive once put it plainly to CBS[34]:

Insiders having an edge to the market is a good thing.

Shayne Coplan, chief executive of Polymarket

For accuracy, in a narrow sense, he is right: the price moved to the truth before the public knew it. It is not prediction, though, and it is no use to an ordinary trader, who is the one on the other side of each of those trades.

The company now says it referred the Maduro trader to the Justice Department itself. In Congress, one bill would bar officials and their staff from trading on government decisions[35], and another would ban markets on elections, sports and military action altogether[36].

Resolution: when the truth is decided by a vote

A share pays a dollar if the event happens, but someone has to say whether it happened. On Polymarket that someone is usually UMA’s optimistic oracle. A proposer posts an outcome with a bond, typically $750, and from there it is a chain of deadlines[37]:

a proposaltwo hours to dispute ita second disputea vote of UMA holders, two days

Most markets resolve in two hours without a dispute. The ones that don’t are where prediction becomes a governance problem.

In March 2025 a $7 million market on whether Ukraine would agree to a minerals deal with the US before April resolved Yes, although Ukraine had not agreed; a single large token holder was accused of swinging the vote, and Polymarket called the situation “unprecedented” but paid no refunds, since “this wasn’t a market failure”[38].

That summer a market with hundreds of millions of dollars traded asked whether Volodymyr Zelensky would be photographed in a suit; he appeared at a NATO summit in a black jacket and matching trousers that many outlets called a suit, and UMA’s voters decided it was not[39].

In May 2026 Bloomberg found that nine wallets held about half of the UMA votes cast in Polymarket disputes over three years[40].

The platform has been narrowing the problem: since September 2025 only whitelisted proposers can propose outcomes, which cut disputes by two thirds in the first month[41], crypto price markets now settle from Chainlink prices without a vote, and a new contract system that can combine several oracles is due to switch on for new markets this week[42].

For a trader, resolution risk is a cost like the spread, small on most markets and large on exactly the ones with ambiguous wording.

Wash trading: volume that isn’t

Volume is the number every article about Polymarket leads with, and it is the least reliable number on the platform.

Researchers at Columbia classified nearly 25% of Polymarket’s historical volume as consistent with wash trading, traders buying from and selling to themselves without changing their position; it peaked at nearly 60% of weekly volume in December 2024, and in sports markets it reached 45% of all-time volume[43].

Separately, Paradigm showed in December 2025 that most dashboards counted every trade twice, once for each side[44].

Neither changes the prices this post studies: a wash trade moves money in a circle, not the price. But together with the trap above, they are good reasons not to read a market’s volume as a measure of how much it knows.

What a Polymarket price is worth

As a probability, a Polymarket price on a market people actually trade is about as good as public forecasts get.

A week before the close the prices were calibrated to within 1.8 percentage points on average, and on the markets open for at least a month they removed 70% of the base-rate forecaster’s error a week out and 82% a day out.

When a liquid market says 70%, the honest reading is “about 70%”, give or take a few points.

That is a statement about many markets, not about one. The Senate at 71.5 cents is a single draw, and so was Brazil at 80. A calibrated price is still wrong three times in ten when it says 70, which is what 70 means.

The weak spots are where theory and the data agree they should be. Thin markets are not forecasts at all, and in them longshots are overpriced. Prices months before the close drift toward 50 cents, because waiting costs money and Polymarket pays for the wait on only a handful of markets.

Crypto price markets share one driver and resolve together, so a bad month for Bitcoin looks like miscalibration. Insiders sometimes move prices before the news, which makes the forecast right and the trade unfair. And the dollar arrives only if the oracle agrees that the event happened.

For a trader the useful findings are modest and specific. The mispricing that survives is at the tails: contracts under 5 cents lose a quarter to two fifths of the stake on average, and near-certain favourites earn a fraction of a percent a week, more when they are bought a month or more before the close.

It is about the size of the spread and the fee, which is why it survives. Nothing in the data suggests a general bias that a formula could exploit, and a two-number correction fitted on the past made the future no better.

The people who make money on Polymarket, according to a study of 588 million trades, are mostly the ones who post limit orders and let others cross the spread; the top 1% of profitable accounts took 76.5% of the profits[45].

For an engineer the lessons are the ones from evaluating any model, sharpened by a market’s habits. Accuracy flatters when the classes are unbalanced, and most questions resolve No. Comparisons across horizons are confounded by which questions exist at each horizon, so they have to be made on the same questions.

And a filter on anything that depends on the outcome, a market’s total volume included, manufactures effects that aren’t there. The trap that caught my first sample is the most transferable thing in this post.

References

  1. CNBC (2026). Lula wins Brazil election despite prediction markets favoring Bolsonaro. 26 October 2026. ↩
  2. Semafor (2026). Prediction markets flop on recent elections. 30 October 2026. ↩
  3. Polymarket (2026). pUSD. Documentation. ↩
  4. Polymarket (2026). Prices and the order book. Documentation. ↩
  5. Wolfers, J., & Zitzewitz, E. (2006). Interpreting prediction market prices as probabilities. NBER Working Paper 12200. ↩
  6. Manski, C. F. (2006). Interpreting the predictions of prediction markets. Economics Letters, 91(3), 425–429. ↩
  7. Polymarket (2026). Trading fees. Documentation. ↩
  8. Guo, C., Pleiss, G., Sun, Y., & Weinberger, K. Q. (2017). On calibration of modern neural networks. ICML 2017. ↩
  9. Dimitriadis, T., Gneiting, T., & Jordan, A. I. (2021). Stable reliability diagrams for probabilistic classifiers. PNAS, 118(8), e2016191118. ↩
  10. Brier, G. W. (1950). Verification of forecasts expressed in terms of probability. Monthly Weather Review, 78(1), 1–3. ↩
  11. Murphy, A. H. (1973). A new vector partition of the probability score. Journal of Applied Meteorology, 12(4), 595–600. ↩
  12. 538 (2024). 538’s final forecasts for the 2024 presidential election. Republished by 6abc, 5 November 2024. ↩
  13. The Wall Street Journal (2024). The Trump betting whale speaks out: “I have absolutely no political agenda”. Republished by To Vima, 1 November 2024. ↩
  14. CBS News (2025). French whale made over $80 million on Polymarket betting on Trump’s election win. 60 Minutes, 30 November 2025. ↩
  15. Horvitz, D. G., & Thompson, D. J. (1952). A generalization of sampling without replacement from a finite universe. Journal of the American Statistical Association, 47(260), 663–685. ↩
  16. Walker, A. (2026). The price is (almost) right: a calibration analysis of Polymarket pricing data. Senior thesis, Princeton University. ↩
  17. Polymarket (2026). How accurate is Polymarket? Accessed 31 October 2026. ↩
  18. Griffith, R. M. (1949). Odds adjustments by American horse-race bettors. American Journal of Psychology, 62(2), 290–294. ↩
  19. Snowberg, E., & Wolfers, J. (2010). Explaining the favorite–long shot bias: is it risk-love or misperceptions? Journal of Political Economy, 118(4), 723–746. ↩
  20. Bürgi, C., Deng, W., & Whelan, K. (2026). Makers and takers: the economics of the Kalshi prediction market. GWU Center for Economic Research Working Paper 2026-001. ↩
  21. Cardozo, M., & Rivero-Wildemauwe, J. I. (2026). The favorite-longshot bias in prediction markets: evidence from Polymarket. arXiv:2609.12878. ↩
  22. Federal Reserve Bank of St. Louis (2026). 3-month Treasury bill secondary market rate (DTB3). FRED. ↩
  23. Page, L., & Clemen, R. T. (2013). Do prediction markets produce well-calibrated probability forecasts? The Economic Journal, 123(568), 491–513. ↩
  24. Polymarket (2026). Holding rewards. Help center, 1 June 2026. ↩
  25. Niculescu-Mizil, A., & Caruana, R. (2005). Predicting good probabilities with supervised learning. ICML 2005. ↩
  26. Le, N. A. (2026). Decomposing crowd wisdom: domain-specific calibration dynamics in prediction markets. arXiv:2602.19520. ↩
  27. Decrypt (2025). Nobel Peace Prize organizers investigating Polymarket insider trades. 10 October 2025. ↩
  28. Bloomberg (2025). Unusual Polymarket bets on Peace Prize winner spur probe. Via Bloomberg Law, 10 October 2025. ↩
  29. Reuters (2026). Nobel committee says Peace Prize winner likely revealed early by digital spying. Via CP24, 30 January 2026. ↩
  30. U.S. Department of Justice (2026). U.S. soldier charged with using classified information to profit from prediction market bets. 23 April 2026. ↩
  31. Commodity Futures Trading Commission (2026). Press release 9217-26. 23 April 2026. ↩
  32. The Block (2026). Fresh accounts netted $1 million on Polymarket hours before US airstrikes on Iran. 28 February 2026. ↩
  33. CBS News (2026). Google employee charged with insider trading on Polymarket. 28 May 2026. ↩
  34. CBS News (2026). Polymarket trader made $400,000 on Maduro’s capture. 6 January 2026. ↩
  35. Torres, R. (2026). Rep. Ritchie Torres introduces legislation to crack down on insider trading on prediction markets. Press release. ↩
  36. Raskin, J. (2026). Raskin, Merkley legislation would ban prediction market gambling on elections, sports, war and government activity. Press release, 26 March 2026. ↩
  37. Polymarket (2026). Resolution. Documentation. ↩
  38. The Block (2025). Polymarket says governance attack by UMA whale to hijack a bet’s resolution is “unprecedented”. 26 March 2025. ↩
  39. DL News (2025). Polymarket bettors clash over Zelenskyy suit dispute. 30 June 2025. ↩
  40. Bloomberg (2026). Las disputas en Polymarket exponen el lado oscuro de los mercados de predicción. Via Bloomberg Línea, 29 May 2026. ↩
  41. UMA (2025). Managed proposers: an update. Blog, 8 November 2025. ↩
  42. The Block (2026). Polymarket rolls out V2 upgrade. 6 October 2026. ↩
  43. Sirolly, A., Ma, H., Kanoria, Y., & Sethi, R. (2025). Network-based detection of wash trading. SSRN Working Paper 5714122. ↩
  44. Paradigm (2025). Polymarket volume is being double-counted. 8 December 2025. ↩
  45. Akey, P., Grégoire, V., Harvie, N., & Martineau, C. (2026). Who wins and who loses in prediction markets? Evidence from Polymarket. CEPR Discussion Paper 21615. ↩
]]> Eva Koroleva AI crypto trading bots: what machine learning and LLMs can and can’t do 2026-11-03T09:00:00+00:00 2026-11-03T00:00:00+00:00 https://blog.korolevalab.com/blog/2026/ai-crypto-trading

In short

Machine learning finds a real edge in crypto prices, but a small one: gradient boosting calls the direction of Bitcoin’s next hour right 53.9% of the time, and an ordinary trader’s fees eat all of it.

What it forecasts well is volatility, which sizes positions and keeps the risk of a strategy predictable, and that is where professional firms use it.

Language models have shown no trading skill in any test long and clean enough to tell, and in backtests on history they remember what came next. An agent that holds keys is a security problem before it is a trading strategy.

63%hit rate an hourly Bitcoin strategy needs to cover a 10 bp round trip
1.56Sharpe ratio of the best of 20,000 random strategies that know nothing
60.0%right on Bitcoin's monthly direction a year or more before the models' release; 46.8% after

In October 2025 a small lab called Nof1 gave six of the strongest language models $10,000 each and let them trade crypto perpetual futures on Hyperliquid with real money.

Every two or three minutes each model received the same prompt with prices and indicators for six coins and answered with a trade: the coin, long or short, the size, the leverage and a plan for getting out[1].

Seventeen days later Qwen3-Max was up 23% and DeepSeek V3.1 up 5%, while Claude Sonnet 4.5, Grok 4, Gemini 2.5 Pro and GPT-5 had lost between a third and two thirds of their money. Bitcoin ended the period where it began.

Account values of six language models trading crypto perpetual futures with $$10,000 each in Alpha Arena, season 1, 18 October to 3 November 2025, on a log scale, with Bitcoin held as a reference (-0.03%). Final returns: Qwen3-Max +23%, DeepSeek V3.1 +5%, Claude Sonnet 4.5 -33%, Grok 4 -48%, Gemini 2.5 Pro -55%, GPT-5 -63%. Grey bands show where a trader with no skill and the same volatility would end up; every model ends inside them Account values of six language models trading crypto perpetual futures with $$10,000 each in Alpha Arena, season 1, 18 October to 3 November 2025, on a log scale, with Bitcoin held as a reference (-0.03%). Final returns: Qwen3-Max +23%, DeepSeek V3.1 +5%, Claude Sonnet 4.5 -33%, Grok 4 -48%, Gemini 2.5 Pro -55%, GPT-5 -63%. Grey bands show where a trader with no skill and the same volatility would end up; every model ends inside them

The results went around social media as a ranking of which model “can trade”. The chart says something more modest.

The grey bands show where a trader with no skill at all, taking the same risk as the models, would end up: two times in three inside the darker band, nineteen times in twenty inside the lighter one.

Every model finished inside the bands, and the winner’s lead is a third of what luck alone would typically give.

Nof1 said as much before the season ended: “We do not expect any of the models to do well, and early successes may be the result of luck.”

Its founder summed up the season with a sentence that is more useful than the leaderboard: “LLMs don’t really handle numerical time series data very well, but that’s all the context we gave them”[2].

Meanwhile, AI has become part of how people trade crypto. In a survey that the exchange MEXC ran among its users in September 2026, 84.8% had used AI for research or trading in the previous six months, and 13% had tried an agent connected to their exchange account[3].

Binance, OKX and Kraken now ship toolkits that let a language model place orders, and regulators keep charging schemes that sold AI trading bots where, in the words of one SEC release, “there were no AI trading bots”[4].

This post tries to separate what models can do in crypto trading from what they are sold for.

I ran the experiments myself on six years of Binance’s public data: gradient boosting that predicts the next hour, a volatility model, almost 800 backtested strategies and 20,000 random ones, two foundation models for time series, and seven open language models asked about Bitcoin’s past.

For ML engineers it explains what makes markets unlike other prediction problems; for traders, what the models actually learn and where the numbers in their marketing come from.

What a model sees in crypto prices

A trading model learns from returns, the relative changes of price from one bar to the next. An hourly return of Bitcoin is \(r_t = \ln(P_t / P_{t-1})\), where \(P_t\) is the price at the close of hour \(t\); for small moves it is the percentage change.

Two facts about these numbers shape everything that follows, and both are easy to see in Binance’s public data.

The first is the shape of their distribution. Measured in standard deviations, most hours are quiet, and a few are wild in a way a normal distribution does not allow:

Histogram of Bitcoin's hourly returns in standard deviations on a log scale, against a normal law. 96 of 51,057 hours moved more than six standard deviations; a normal law expects 0.0001 Histogram of Bitcoin's hourly returns in standard deviations on a log scale, against a normal law. 96 of 51,057 hours moved more than six standard deviations; a normal law expects 0.0001

From January 2021 to the end of October 2026 Bitcoin had 51,057 hours, and 96 of them moved more than six standard deviations. Under a normal law with the same mean and spread, a six-sigma hour happens about once in 58,000 years. Ether and Solana look the same, with 108 and 109 such hours.

For a model this means that a handful of hours carry a large share of the total movement, and a loss function that squares errors is pulled around by them. It also means that any test period without a crash is missing the hours that decide whether a leveraged strategy survives.

The second fact is about memory. The autocorrelation of a series at lag \(k\) is its correlation with itself \(k\) steps earlier. For returns it is close to zero at every lag; for the size of the returns, their absolute value, it is not:

Autocorrelation of Bitcoin's hourly returns stays within noise at every lag up to a week (lag 1: -0.005), while that of absolute returns starts at 0.29 and is still 0.16 after a week, with bumps at whole days Autocorrelation of Bitcoin's hourly returns stays within noise at every lag up to a week (lag 1: -0.005), while that of absolute returns starts at 0.29 and is still 0.16 after a week, with bumps at whole days

The direction of the last hour says almost nothing about the direction of the next one, but a violent hour is likely to be followed by another violent hour, and that memory lasts for days.

The bumps at 24, 48 and 72 hours are the daily rhythm of trading. The same hours are busy every day: in the hour from 14:00 UTC, around the opening of stock trading in New York, Bitcoin moves 56 bp on average, against 29 bp in the hour from 04:00 (a basis point, bp, is 0.01%).

These two curves are the plan of this post in miniature. Direction is close to unpredictable, so a model that predicts it will be right only a little more often than a coin, and the question is whether that little survives the costs of trading.

Volatility is predictable, and that is where machine learning earns its keep without anyone writing about it.

Predicting the next hour

The experiment is the one a newcomer to quantitative trading usually runs first, done carefully.

Every hour, a model sees what a trader’s screen shows at the close of that hour for Bitcoin, Ether and Solana: returns over the last 1, 2, 4, 8, 24, 72 and 168 hours, Bitcoin’s returns as a market factor, volatility measured from 5-minute candles, volume against its weekly norm, the share of volume bought by takers (traders who cross the spread and pay for immediacy), the range of the last hour, the hour of the day, the day of the week and which coin it is.

Twenty features in all.

It predicts the probability that the next hour closes higher. The model is gradient boosting, LightGBM with shallow trees, the default choice for tabular data and a strong one.

The test has to respect time. Shuffling hours into random training and test sets would let the model learn from the future, so the evaluation is walk-forward:

Walk-forward testing: for each of 58 months from January 2022 to October 2026, a new model is trained on all hours before the month, less a gap of a day, and predicts that month only. The months it predicts are joined into one out-of-sample record Walk-forward testing: for each of 58 months from January 2022 to October 2026, a new model is trained on all hours before the month, less a gap of a day, and predicts that month only. The months it predicts are joined into one out-of-sample record

Every month from January 2022 to October 2026 gets its own model, trained on every hour before that month except the last day, which is kept out so that no training target overlaps the test. The 58 months of predictions, 42,142 hours per coin, form one long out-of-sample record that no model saw while it was trained.

The result is better than I expected. The model calls the direction of Bitcoin’s next hour correctly 53.9% of the time, against 50.5% for always saying “up”; for Ether 53.8%, for Solana 52.9%.

The area under the ROC curve is 0.555 for Bitcoin. Over 42,142 hours a hit rate this far from a coin toss is not luck: its standard error is a quarter of a percentage point.

Trap

A result like this deserves suspicion before celebration. The usual cause of a good backtest is a leak, a feature that quietly contains the future. Here the hourly closes rebuilt from 5-minute candles match Binance’s own hourly candles to the cent, and the pattern the model found, described next, survives when the last five minutes before each forecast are skipped.

The feature that carries most of the signal is the simplest: the return of the last couple of hours, with a negative sign. Crypto hourly returns tend to be partly taken back in the next hour.

The rank correlation between consecutive hourly returns of Bitcoin is −0.06 over the test period, small but steady, and it shows plainly in how often an hour moves the same way as the hour before:

Share of hours that move in the same direction as the previous hour, by year, for BTC, ETH and SOL. Every bar is below 50%, lowest in 2022 and 2023: BTC 2021 48.6%, 2022 46.0%, 2023 45.9%, 2024 46.7%, 2025 48.7%, 2026 47.7%; ETH 2021 48.5%, 2022 46.8%, 2023 46.1%, 2024 46.3%, 2025 48.2%, 2026 47.7%; SOL 2021 49.2%, 2022 48.0%, 2023 46.6%, 2024 46.9%, 2025 49.3%, 2026 47.7% Share of hours that move in the same direction as the previous hour, by year, for BTC, ETH and SOL. Every bar is below 50%, lowest in 2022 and 2023: BTC 2021 48.6%, 2022 46.0%, 2023 45.9%, 2024 46.7%, 2025 48.7%, 2026 47.7%; ETH 2021 48.5%, 2022 46.8%, 2023 46.1%, 2024 46.3%, 2025 48.2%, 2026 47.7%; SOL 2021 49.2%, 2022 48.0%, 2023 46.6%, 2024 46.9%, 2025 49.3%, 2026 47.7%

A mean-reverting wiggle like this is what liquidity providers are paid for. When an impatient trader pushes the price by buying a lot at once, someone has to sell to them, and that someone expects the price to drift back. The effect was strongest in 2022 and 2023 and weaker in 2025 and 2026.

Where the edge goes

A model that is right 53.9% of the time sounds like a money machine, and this is the moment to look at what trading costs.

Every trade pays a fee, and every market order also pays the spread and walks up the order book. Suppose a strategy opens and closes a position on every bar, and the average absolute move over the bar is \(m\).

It earns \(+m\) when right and \(-m\) when wrong, and pays the round-trip cost \(c\) every time. With hit rate \(p\) the expected profit per bar is

\[\mathbb{E}[\text{profit}] = p\,m - (1 - p)\,m - c = (2p - 1)\,m - c,\]

which is positive only if

\[p > \frac{1}{2} + \frac{c}{2m}.\]

The ratio of cost to typical move decides everything. Bitcoin moves 38 bp in an average hour, and a round trip with taker fees on Binance’s futures costs about 10 bp, so an hourly strategy needs to be right 63% of the time. On 5-minute bars, with an average move of 11 bp, it needs 94%:

The hit rate a strategy trading every bar needs to cover its costs. 5 min: average move 11 bp, hit rate needed 94% at 10 bp; 15 min: average move 19 bp, hit rate needed 76% at 10 bp; 1 hour: average move 38 bp, hit rate needed 63% at 10 bp; 4 hours: average move 78 bp, hit rate needed 56% at 10 bp; 1 day: average move 203 bp, hit rate needed 52% at 10 bp The hit rate a strategy trading every bar needs to cover its costs. 5 min: average move 11 bp, hit rate needed 94% at 10 bp; 15 min: average move 19 bp, hit rate needed 76% at 10 bp; 1 hour: average move 38 bp, hit rate needed 63% at 10 bp; 4 hours: average move 78 bp, hit rate needed 56% at 10 bp; 1 day: average move 203 bp, hit rate needed 52% at 10 bp

The fee is only part of the cost. A market order also crosses the spread and, if it is large, eats into the order book.

I took six snapshots of the books of Binance’s perpetual futures, five minutes apart, and measured what a market buy of a given size costs against the mid price, for contracts ranked by their daily turnover:

Cost of a market buy against the mid price on Binance perpetual futures, by order size from $$1,000 to $$1 million, for contracts ranked by daily turnover, log scales. #1 BTC: $$100K costs 0.01 bp; #2 ETH: $$100K costs 0.02 bp; #5 HYPE: $$100K costs 3.22 bp; #10 XRP: $$100K costs 1.07 bp; #25 AVAX: $$100K costs 6.10 bp; #50 PENGU: $$100K costs 10.56 bp; #100 1000FLOKI: $$100K costs 31.20 bp; #200 MET: $$100K costs 86.88 bp Cost of a market buy against the mid price on Binance perpetual futures, by order size from $$1,000 to $$1 million, for contracts ranked by daily turnover, log scales. #1 BTC: $$100K costs 0.01 bp; #2 ETH: $$100K costs 0.02 bp; #5 HYPE: $$100K costs 3.22 bp; #10 XRP: $$100K costs 1.07 bp; #25 AVAX: $$100K costs 6.10 bp; #50 PENGU: $$100K costs 10.56 bp; #100 1000FLOKI: $$100K costs 31.20 bp; #200 MET: $$100K costs 86.88 bp

For Bitcoin and Ether the book is so deep that even a million-dollar order costs a fraction of a basis point: 0.18 bp for Bitcoin, 1.3 bp for Ether. Further down the list the picture changes quickly.

On the 25th most traded contract, a $100,000 order costs 6 bp before any fee; on the 50th, 11 bp; on the 200th, 87 bp, and the visible book cannot fill a million at all.

These snapshots were taken on a Saturday, a day and a half after Binance disclosed a $6.3 billion hack of one of its hot wallets and suspended withdrawals[5], so the books of small contracts may be thinner than on an ordinary weekday.

The point stands on any day: the smaller the coin, the more of a model’s edge goes to the order book.

Now the model’s trades. The plain strategy is long when the model gives the next hour more than even odds to rise and short otherwise, and it pays a fee whenever it switches sides:

Value of one dollar traded on the model's hourly forecasts for Bitcoin from January 2022 to October 2026, log scale: without fees it ends at 9.9 dollars, holding Bitcoin at 1.65, with 2 bp per side at 0.036, with 5 bp per side at practically zero Value of one dollar traded on the model's hourly forecasts for Bitcoin from January 2022 to October 2026, log scale: without fees it ends at 9.9 dollars, holding Bitcoin at 1.65, with 2 bp per side at 0.036, with 5 bp per side at practically zero

Without fees, a dollar traded on Bitcoin’s forecasts from January 2022 grows to about ten. With 2 bp per side, the fee of a market maker on futures, it loses most of its value; with 5 bp, a taker’s fee, it is gone.

The model switches sides 16 times a day, and a 53.9% hit rate does not pay for that.

It is also worse than 53.9% suggests. When the model is right, the hour moves 31.5 bp on average; when it is wrong, 35.6 bp. It is right more often on the small moves and wrong more often on the large ones, so the hit rate overstates the profit, a pattern that every strategy built on small reversals shares.

The obvious fix is to trade less: keep the current position unless the model’s probability crosses a band around one half, say above 0.55 or below 0.45. A wider band means fewer trades and fewer fees, but also acting on fewer of the model’s opinions:

Return per year against how often the strategy changes side, for four fee levels. BTC 16.0 changes a day: no fees 48%, 5 bp -244%; BTC 13.6 changes a day: no fees 42%, 5 bp -205%; BTC 11.3 changes a day: no fees 50%, 5 bp -155%; BTC 9.3 changes a day: no fees 44%, 5 bp -125%; BTC 6.0 changes a day: no fees 57%, 5 bp -53%; BTC 2.6 changes a day: no fees -5%, 5 bp -52%; SOL 15.8 changes a day: no fees 97%, 5 bp -191%; SOL 12.5 changes a day: no fees 117%, 5 bp -112%; SOL 9.8 changes a day: no fees 126%, 5 bp -52%; SOL 7.6 changes a day: no fees 93%, 5 bp -45%; SOL 4.3 changes a day: no fees 63%, 5 bp -15%; SOL 1.4 changes a day: no fees 16%, 5 bp -9% Return per year against how often the strategy changes side, for four fee levels. BTC 16.0 changes a day: no fees 48%, 5 bp -244%; BTC 13.6 changes a day: no fees 42%, 5 bp -205%; BTC 11.3 changes a day: no fees 50%, 5 bp -155%; BTC 9.3 changes a day: no fees 44%, 5 bp -125%; BTC 6.0 changes a day: no fees 57%, 5 bp -53%; BTC 2.6 changes a day: no fees -5%, 5 bp -52%; SOL 15.8 changes a day: no fees 97%, 5 bp -191%; SOL 12.5 changes a day: no fees 117%, 5 bp -112%; SOL 9.8 changes a day: no fees 126%, 5 bp -52%; SOL 7.6 changes a day: no fees 93%, 5 bp -45%; SOL 4.3 changes a day: no fees 63%, 5 bp -15%; SOL 1.4 changes a day: no fees 16%, 5 bp -9%

With a band of ±0.05 the Bitcoin strategy trades 6 times a day instead of 16 and still earns 57% a year before costs (here and below, a year’s hourly returns added up, without compounding).

At 1 bp per side 35% remains, at 2 bp 13%, at 5 bp it loses 53%. A wider band loses the edge before it saves enough on fees. Solana, with larger moves, keeps some profit at 2 bp per side, and none at 5.

For traders

The edge is real, and it belongs to whoever pays the least for trading. A market maker who earns a rebate for resting orders, sits next to the exchange and pays a fraction of a basis point can harvest it; a retail trader paying 5 bp per side through an API cannot, however good the model.

This backtest is also generous: it assumes trades at the exact closing price of the hour, with no delay and no slippage beyond the fee. The real numbers would be worse.

What can be predicted: volatility

The other curve of the autocorrelation chart, the one that stayed high, is where forecasting works. Realized variance is the sum of squared 5-minute returns over a day,

\[RV_t = \sum_{i=1}^{288} r_{t,i}^2,\]

a direct measurement of how much the price moved that day. The HAR model of Corsi (2009) forecasts tomorrow’s value from three averages of it: yesterday’s, the last week’s and the last month’s,

\[\ln RV_{t+1} = \beta_0 + \beta_d \ln RV_t + \beta_w \ln RV^{(w)}_t + \beta_m \ln RV^{(m)}_t + \varepsilon_{t+1}.\]

The three terms stand for traders with different horizons, who react to yesterday, to the last week and to the last month. It is a linear regression with four coefficients, and it is hard to beat[6].

I fitted it the same way as the direction model, refitting every month on all earlier days, and next to it fitted the same regression on returns: tomorrow’s return from yesterday’s, the last week’s and the last month’s.

The score is the out-of-sample \(R^2\), the share of the variation of the target that the forecast explains, against the simplest forecast (the historical mean for volatility, zero for returns):

Left: Bitcoin's realized daily volatility since July 2025 and the HAR forecast made a day earlier, which follows it closely. Right: out-of-sample R squared, volatility against return: BTC 0.54 against -0.003; ETH 0.57 against -0.006; SOL 0.63 against -0.008 Left: Bitcoin's realized daily volatility since July 2025 and the HAR forecast made a day earlier, which follows it closely. Right: out-of-sample R squared, volatility against return: BTC 0.54 against -0.003; ETH 0.57 against -0.006; SOL 0.63 against -0.008

For volatility \(R^2\) is 0.54 for Bitcoin, 0.57 for Ether and 0.63 for Solana; for returns it is below zero for all three, a forecast worse than saying “zero” every day.

A volatility forecast is not a trading signal by itself, but it decides how large a position should be. A position scaled by the inverse of forecast volatility keeps its risk roughly constant, small when the market is wild and larger when it is calm:

Volatility over 30 days of holding Solana, swinging from about 40 to over 150 percent a year, and of a position sized by the forecast, which stays near the target of 50 percent. Realized volatility 93% against 48%, largest drawdown 95% against 72% Volatility over 30 days of holding Solana, swinging from about 40 to over 150 percent a year, and of a position sized by the forecast, which stays near the target of 50 percent. Realized volatility 93% against 48%, largest drawdown 95% against 72%

Holding Solana from 2022 to October 2026 meant a realized volatility of 93% a year and a worst fall of 94.5% from a peak. Scaling the position to a target of 50% held the realized volatility at 48% and cut the worst fall to 72%.

For Bitcoin it raised the Sharpe ratio, the return per unit of volatility, from 0.21 to 0.31. It does not make a bad asset good, but it makes the risk of a strategy predictable, and in a market with six-sigma hours that matters more than a percent of hit rate.

This is the part of machine learning that professional crypto firms use every day and rarely write about: forecasting volatility and liquidity, sizing positions, setting quotes, managing risk. The models are less glamorous than price prediction and much more useful.

The backtest that flatters

Every number so far comes from data the models never saw. Most strategies offered to traders, including many sold as AI, come from the opposite procedure: try many ideas on the same history and keep the one that looks best.

Machine learning engineers know this as overfitting to the validation set; in finance it has its own literature and its own name, backtest overfitting.

To see its size, I ran a grid of the rules every trading textbook contains: moving-average crossovers, breakouts of a price channel, and two kinds of mean reversion (RSI and Bollinger bands), each with a range of parameters, on the hourly candles of twelve large coins.

Sixty-six rules on twelve coins make 792 strategies. Each was scored by its Sharpe ratio, the average return divided by its standard deviation and scaled to a year, after a fee of 5 bp per side, once on 2021-2023 and once on 2024 to October 2026:

Scatter of 792 strategies, Sharpe ratio in 2021-2023 against 2024-2026. The cloud is flatter than the diagonal; the 20 best in the first period, highlighted, average 1.56 then and 0.19 later. Rank correlation 0.45 Scatter of 792 strategies, Sharpe ratio in 2021-2023 against 2024-2026. The cloud is flatter than the diagonal; the 20 best in the first period, highlighted, average 1.56 then and 0.19 later. Rank correlation 0.45

The ranking in the first period was not useless: its rank correlation with the second is 0.45, mostly because it identifies the rules that trade too often and lose to fees in any period. Mean reversion with RSI and Bollinger bands lost money in both. The winners are another matter.

The 20 best strategies of 2021-2023 kept about an eighth of what the backtest promised: an average Sharpe ratio of 0.19 from 2024 on, against 1.56, and 6 of the 20 lost money.

The second experiment removes any doubt about what is skill. Each of 20,000 strategies holds a random position, long or short, and flips it at random about once a day; none of them knows anything. Among \(N\) of them, the best one still looks good, and the more of them one tries, the better the best looks:

Best Sharpe ratio among N random strategies against N on a log scale, with the expected maximum of Bailey and Lopez de Prado. N=1: -0.55; N=2: -0.23; N=5: 0.11; N=10: 0.32; N=20: 0.51; N=50: 0.71; N=100: 0.84; N=200: 0.95; N=500: 1.10; N=1000: 1.20; N=2000: 1.35; N=5000: 1.44; N=10000: 1.48; N=20000: 1.56 Best Sharpe ratio among N random strategies against N on a log scale, with the expected maximum of Bailey and Lopez de Prado. N=1: -0.55; N=2: -0.23; N=5: 0.11; N=10: 0.32; N=20: 0.51; N=50: 0.71; N=100: 0.84; N=200: 0.95; N=500: 1.10; N=1000: 1.20; N=2000: 1.35; N=5000: 1.44; N=10000: 1.48; N=20000: 1.56

Bailey and López de Prado (2014) gave the expected best Sharpe ratio among \(N\) independent trials that have no skill:

\[\mathbb{E}\big[\max_N \widehat{SR}\big] \approx \mathbb{E}\big[\widehat{SR}\big] + \sqrt{V}\left((1-\gamma)\,\Phi^{-1}\!\Big(1 - \frac{1}{N}\Big) + \gamma\,\Phi^{-1}\!\Big(1 - \frac{1}{N e}\Big)\right).\]

Here \(\mathbb{E}[\widehat{SR}]\) and \(V\) are the mean and the variance of the estimated Sharpe ratios across the trials, \(\Phi^{-1}\) is the inverse of the standard normal distribution function, \(\gamma \approx 0.5772\) is the Euler-Mascheroni constant and \(e\) is Euler’s number. The formula is the expected maximum of \(N\) normal draws, and it grows like \(\sqrt{2 \ln N}\): slowly, but without limit.

Here the average random strategy has a Sharpe ratio of −0.55, because it pays fees for nothing. The best of 10 has 0.32, the best of 1,000 has 1.20, and the best of 20,000 has 1.56, as good as the 20 best real rules looked in their own past. The formula, with the trials’ average added, follows the experiment closely.

Their deflated Sharpe ratio turns this into a test: it asks how likely the best strategy’s true Sharpe ratio is to be above zero, given how many strategies were tried and how long and fat-tailed the record is[7].

Applied to the best rule of my grid, the in-sample Sharpe ratio of 1.81 is below the 2.04 that the best of 792 trials would be expected to reach by luck, given how much their results differ, and the probability that its true Sharpe ratio is above zero comes out at 35%.

That rule, a crossover of 12- and 24-hour averages on Dogecoin, did well afterwards anyway, with a Sharpe ratio of 1.21. A deflated Sharpe ratio is a probability, not a verdict; what it says is that the backtest alone was not evidence.

The same arithmetic applies to every published result. Harvey, Liu and Zhu counted at least 316 factors that academic papers had tested to explain stock returns and concluded that, after so much searching, a new one should have a t-statistic above 3 rather than the usual 2 to be taken seriously[8].

A seller of a trading bot who shows the best of their backtests is showing the maximum of \(N\) draws, and rarely says what \(N\) was.

Language models as traders

Short records have a second problem besides selection: they are too short to tell skill from luck even for one strategy. Under the simplest assumptions, the standard error of an annualized Sharpe ratio estimated from \(D\) days of crypto trading is about

\[\mathrm{SE}\big(\widehat{SR}\big) \approx \sqrt{\frac{365}{D}},\]

so 17 days of Alpha Arena give an uncertainty of about 4.6. A model would have needed an annualized Sharpe ratio above 9 to stand out from zero.

Trading more often does not help: the precision is set by calendar time, not by the number of decisions[9]. To show that a strategy has a Sharpe ratio of 1, the usual level of a good one, takes about four years of daily returns.

The academic benchmarks of LLM traders run into the same wall. AI-Trader tested six models on ten coins for two weeks of November 2025; five of the six did worse than an index of the same coins[10].

In StockBench, thirteen models traded twenty Dow stocks for four months and returned between −2.8% and +2.5%, against +0.4% for simply holding the stocks; in a second test over the falling market of early 2025, every agent did worse than holding[11].

LiveTradeBench found that a model’s rank on general benchmarks says nothing about its returns: the rank correlation was 0.05[12].

The largest record I know of is the least flattering. A company that runs two fleets of LLM agents published six months of their trades; on Hyperliquid, 500 to 600 user-created agents lost money, won 41% of their round trips against 50% for comparable human retail traders, and used a median leverage of 5 regardless of how volatile the market was.

Liquidations were 3.2% of the closed positions and 74% of the gross loss[13]. It is the operator’s own paper, which makes the result harder, not easier, to dismiss.

Longer tests are possible only on history, and history brings in a problem specific to language models. A model trained on text up to some date has read what happened before it. An LLM asked in 2026 whether to buy Bitcoin on 1 March 2020 may simply remember what came next.

To see how much, I asked seven open models released between May 2024 and February 2026 about Bitcoin, from January 2018 to October 2026.

Each model answered three kinds of questions: what Bitcoin closed at on the 1st and the 15th of every month, whether it went up or down over a given month, and whether it went up or down over a given week.

For the up-or-down questions I did not let the models write an answer; I read the probabilities they gave to “up” and “down” as the first word, so that a model that would rather explain that it cannot know still had to lean one way. The models ran on Kaggle’s GPUs, in 16-bit precision and without any tools or web search.

The price answers show where each model’s knowledge ends:

Bitcoin's real closing price from 2018 to 2026 and four open models' answers for the 1st and 15th of each month. Each model's answers follow the price closely up to some point before its release and then stop, turn into refusals, zeros or a constant guess. Mistral 7B v0.3, released 2024-05: median error 5% on dates a year or more before its release; Qwen2.5 7B, released 2024-09: median error 9% on dates a year or more before its release; Olmo 3 7B, released 2025-11: median error 18% on dates a year or more before its release; Qwen3.5 9B, released 2026-02: median error 10% on dates a year or more before its release Bitcoin's real closing price from 2018 to 2026 and four open models' answers for the 1st and 15th of each month. Each model's answers follow the price closely up to some point before its release and then stop, turn into refusals, zeros or a constant guess. Mistral 7B v0.3, released 2024-05: median error 5% on dates a year or more before its release; Qwen2.5 7B, released 2024-09: median error 9% on dates a year or more before its release; Olmo 3 7B, released 2025-11: median error 18% on dates a year or more before its release; Qwen3.5 9B, released 2026-02: median error 10% on dates a year or more before its release

Mistral 7B, released in May 2024, named Bitcoin’s closing price for dates in 2018 to 2022 with a median error of 3 to 5%; for every date after its release it declined to give a number.

Qwen3.5 9B, released in February 2026, was off by 4 to 5% on dates in 2023 and 2024, still close in January and February 2025, and for every date from March 2025 on it answered “0”: its knowledge ends a year before its release.

Olmo 3 7B had a median error of 9 to 10% for 2023 and 2024 and explained, in one of its refusals, that it had no data “beyond my knowledge cutoff in June 2023”, which its own answers contradict.

None of these models was trained to know prices. They picked them up from the internet the way they picked up everything else.

Direction is what a trading backtest needs, and the models remember it too:

Share of months for which each model says correctly whether Bitcoin went up or down, for months that ended a year or more before the model's release and for months after it: Mistral 7B v0.3 57% of 65 months before, 50% of 28 after; Qwen2.5 7B 59% of 69 months before, 46% of 24 after; OLMo 2 7B 67% of 72 months before, 52% of 21 after; Phi-4-mini 53% of 74 months before, 47% of 19 after; Qwen3 8B 58% of 76 months before, 41% of 17 after; Olmo 3 7B 66% of 83 months before, 50% of 10 after; Qwen3.5 9B 59% of 86 months before, 29% of 7 after Share of months for which each model says correctly whether Bitcoin went up or down, for months that ended a year or more before the model's release and for months after it: Mistral 7B v0.3 57% of 65 months before, 50% of 28 after; Qwen2.5 7B 59% of 69 months before, 46% of 24 after; OLMo 2 7B 67% of 72 months before, 52% of 21 after; Phi-4-mini 53% of 74 months before, 47% of 19 after; Qwen3 8B 58% of 76 months before, 41% of 17 after; Olmo 3 7B 66% of 83 months before, 50% of 10 after; Qwen3.5 9B 59% of 86 months before, 29% of 7 after

Over the months that ended at least a year before a model’s release, the seven models together called the direction right 60.0% of the time, in 525 answers; over the months after their release, 46.8% of 126.

Each model’s answers are split at its own median probability, so that a model that says “up” to everything is judged by how it orders the months.

The newest models have only a few months after their release, so their grey bars are noisy; the pooled numbers are not. For weeks the effect is smaller, 53.3% against 49.0%, probably because far less is written about single weeks than about months.

Now suppose someone builds a trading strategy on these answers and backtests it on 2018-2026, as one would with any signal:

A dollar traded monthly on OLMo 2 7B's answers grows to 245 dollars before the model's release in 2024-12, against 7 for holding Bitcoin, and then goes to 198, a change of -19% after the release A dollar traded monthly on OLMo 2 7B's answers grows to 245 dollars before the model's release in 2024-12, against 7 for holding Bitcoin, and then goes to 198, a change of -19% after the release

Trading OLMo 2’s monthly answers turns a dollar into 245 between January 2018 and November 2024, while holding Bitcoin turns it into 7. After the model’s release in December 2024 the same strategy loses 19% over 22 months. Nothing in the backtest is wrong except the one thing that matters: the model had read the answers.

Across the seven models, the strategy multiplied money by 8 to 245 before each model’s release and by 0.54 to 1.31 after it.

This is not a curiosity. Lopez-Lira, Tang and Zhu asked GPT-4o for daily levels of the S&P 500: before its training cutoff the error was 0.61%, after it 16.87%.

Telling the model to ignore what it knows did not help, and anonymizing companies did not either: from an earnings call with the names removed, it identified Apple, Meta and Microsoft every time[14].

When five published LLM trading agents were rerun on a period after their model’s cutoff, with a market return similar to their original period, their returns fell by 50 to 72%[15].

Over twenty years of data instead of a few months, holding the stocks beat the two best-known LLM trading agents in every setup[16].

There are fixes. Point-in-time models are trained only on text available up to a given date, one model per year, so that a backtest for 2015 uses a model that has never read 2016[17].

Their authors found the bias modest for one task, ranking stocks by the tone of the day’s news, which is a reminder that the size of the problem depends on what the model is asked.

For anything that names a coin and a date, the only clean test is the one that runs after the model was trained, live, for long enough.

Where language models help

None of this makes language models useless to a trader. They are weak at the job that is easiest to sell, taking prices in and sending trades out, and good at several jobs around it.

The clearest case is research. Man Group, one of the largest quantitative hedge funds, describes a pipeline of LLM agents in which one proposes hypotheses, another writes the code to test them on the firm’s data and a third runs the statistical tests; in their words, “whether generated by human or AI, every strategy entering live trading must pass identical thresholds”[18].

The thresholds matter because a model that writes a hundred strategies an hour is the fastest way yet to make \(N\) in the overfitting formula large.

Code is the second case. Coinbase says that between 95% and 100% of its code is now written by or with LLMs[19].

Data pipelines, backtests, exchange connectors and monitoring are where a trading team spends most of its time, and a model that writes them faster saves more than one that guesses the next candle. Reading is the third: whitepapers, governance proposals, the code of a smart contract.

Anthropic’s red team found that models could write working exploits for half of 405 smart contracts that had been hacked in the past, in simulation, and found two new vulnerabilities in recently deployed contracts[20]; the same skill reviews a contract before one deposits into it, and attacks it for someone else.

News is the case everyone expects to work. The evidence is thin: a 2026 study scored Bitcoin news with ChatGPT and found that the sentiment improved the fit in sample but added nothing significant out of sample[21].

Professional market makers, for their part, still hire for older skills. A recent job advertisement of Wintermute, one of the largest, asks a machine learning researcher to “develop ML-based alpha generation models using high-frequency order book and market microstructure data” and does not mention language models[22].

Forecasting events, rather than prices, is where language models have come furthest. On ForecastBench, a benchmark of questions about the future, several systems were by July 2026 “statistically indistinguishable from superforecaster-level accuracy”, and on questions taken from prediction markets one ranked above the superforecasters’ median for the first time; the human forecasts it was compared with were collected in 2024[23].

The money that bots have demonstrably made in prediction markets comes from something plainer: buying a set of outcomes that must together pay $1 for less than $1. On Polymarket this kind of arbitrage extracted an estimated $40 million in a year[24].

Then there are foundation models for time series, Transformers pretrained on millions of series in the way language models are pretrained on text.

Kronos was trained on more than 12 billion candles from 45 exchanges and reads open, high, low, close and volume as tokens[25]; Chronos-2 is a general forecaster from Amazon[26].

I found no independent test of them on crypto returns, so I ran one: both models, without fine-tuning, forecast every hour of 2026 for Bitcoin, Ether and Solana from the 512 hours before, and Chronos-2 also forecast daily volatility.

Left: share of hours of 2026 where the direction of the next hour was right, for LightGBM, Chronos-2 and Kronos-small. BTC 53.2%, 49.3%, 52.2%; ETH 52.3%, 49.6%, 50.7%; SOL 53.0%, 49.2%, 49.6%. Right: R squared of tomorrow's log volatility, HAR against Chronos-2: BTC 0.77 against 0.67; ETH 0.83 against 0.64; SOL 0.89 against 0.66 Left: share of hours of 2026 where the direction of the next hour was right, for LightGBM, Chronos-2 and Kronos-small. BTC 53.2%, 49.3%, 52.2%; ETH 52.3%, 49.6%, 50.7%; SOL 53.0%, 49.2%, 49.6%. Right: R squared of tomorrow's log volatility, HAR against Chronos-2: BTC 0.77 against 0.67; ETH 0.83 against 0.64; SOL 0.89 against 0.66

Neither model knows the direction of the next hour. Chronos-2 was right 49.2 to 49.6% of the time, slightly worse than a coin, and said “up” in 66 to 75% of the hours: from a series of hourly returns it learned mostly that the average is a little above zero.

Kronos-small was right 49.6 to 52.2% of the time on the hours it forecast; with 1,812 hours per coin, a deviation of one or two points from 50% is within noise.

The gradient-boosting model with its twenty hand-made features did better on the same months, which is what one would expect from a model trained on this very market.

Volatility tells the same story as before. Chronos-2, which has never been trained on crypto in particular, forecast tomorrow’s log volatility with an \(R^2\) of 0.64 to 0.67 against the average of the previous 512 days; the HAR regression with four coefficients scored 0.77 to 0.89 on the same days.

A large pretrained model learned that volatility clusters, and a model built around that fact still knows it better.

Independent tests on stock returns found the same: zero-shot Chronos and TimesFM had negative out-of-sample \(R^2\) on returns, where gradient boosting was close to zero[27], and on realized volatility only a tiny model beat the logarithmic HAR[28].

Agents with keys

In crypto, “AI agent” first meant something other than a trading system. In 2024 a chatbot called Truth Terminal, which posted on X, received a $50,000 research grant in bitcoin from Marc Andreessen and then promoted a memecoin that someone else had launched[29].

Tokens of other “AI agents” followed, and at their peak around the turn of 2025 they were worth about $16 billion together; most have since fallen by more than 90%[30]. In August 2026 the founder of ai16z, a token once worth $2.4 billion, wrote: “The token is dead. Completely”[31].

The agents themselves outlived their tokens, and the interesting question about them is not price but security. A program that reads the internet and can move money is a new kind of target, because everything it reads is text, and a language model treats text as a possible instruction.

The first public demonstration was a game. Freysa was an agent with a wallet and one rule, never to send the money; players paid a rising fee per message to try to talk it into doing so.

After 482 attempts one player convinced it that its transfer function was for receiving funds, offered a donation, and took the prize of $47,316[32].

The real incidents are less elegant. In March 2025 someone with access to the dashboard of AIXBT, an agent that posts market commentary, queued two replies that made it send about 55 ETH, some $100,000, from its wallet; the team said it was “not a result of agent manipulation”, which makes it an ordinary access-control failure with an AI attached[33].

In May 2026 an attacker sent Grok’s wallet a membership NFT that widened what the wallet was allowed to transfer, then posted a message in Morse code; Grok decoded it into an instruction to a trading bot, Bankrbot, which sent 3 billion DRB tokens, worth about $150,000 to $200,000[34]. No key was stolen and no contract had a bug.

Memory makes this worse. Researchers at Princeton showed that an attacker can plant a false instruction in an agent’s long-term memory through an ordinary conversation, and the agent later acts on it when other users make legitimate requests. In their summary, “even the strongest models followed poisoned memories more than half the time”[35].

The tools around the model are a supply chain of their own. Of 2,857 skills in one open marketplace for agents, a security firm found 341 malicious, 163 of them dressed as crypto tools and prediction-market bots[36].

An MCP server published on npm asked for a wallet’s private key to “auto-sign” payments and sent it to a remote server with every call[37].

Of 428 third-party routers that resell access to LLM APIs, researchers found nine that injected malicious code, and one that drained ETH from a key planted as bait[38].

Even the news can be an attack: altering headlines with invisible characters on a single day cut a sentiment-trading system’s annual return by up to 17.7 percentage points[39].

Parts of an LLM trading agent. Inputs: prices and order books, news and social posts, on-chain data, and the agent's memory of earlier sessions. All of them reach the language model as text, and the model calls tools and plugins. It proposes an order, which passes through risk checks written in ordinary code (position and loss limits, allowed markets) before it reaches the exchange with an API key that can trade but not withdraw. Attack points: instructions hidden in a post, planted memories, a malicious plugin or tool, and a leaked key Parts of an LLM trading agent. Inputs: prices and order books, news and social posts, on-chain data, and the agent's memory of earlier sessions. All of them reach the language model as text, and the model calls tools and plugins. It proposes an order, which passes through risk checks written in ordinary code (position and loss limits, allowed markets) before it reaches the exchange with an API key that can trade but not withdraw. Attack points: instructions hidden in a post, planted memories, a malicious plugin or tool, and a leaked key

The exchanges’ agent toolkits of 2026 have converged on the same defence. Binance’s lets agents trade from a walled-off sub-account but not move assets to an external wallet[40].

OKX signs orders on the user’s machine, so that “the AI client sends only trading intent”, and offers a mode with simulated funds[41]; Kraken’s command-line tool has a local paper-trading engine[42]. A key that cannot withdraw prevents theft. It does nothing against an agent that loses money by trading.

For engineers

The design that follows from all this is old-fashioned. The model reads untrusted text and proposes; ordinary code decides, with limits on position size, leverage and daily loss that the model cannot change. Memory is per user and is treated as untrusted input too. Tools come from sources one has read. And the model never holds a key that can send funds anywhere.

Bots that sell returns

Every claim in the previous sections can be checked against data. The market for “AI trading bots” mostly cannot, and regulators have been warning about it for years. In January 2024 the US Commodity Futures Trading Commission published an advisory titled “AI Won’t Turn Trading Bots into Money Machines”, with a sentence that could close this post[43]:

AI technology can’t predict the future or sudden market changes.

US Commodity Futures Trading Commission, customer advisory

Its main example was Mirror Trading International, which collected at least 29,421 bitcoin, worth $1.7 billion, from about 23,000 Americans for a trading “bot”[44]. In the advisory’s words, “in reality, very little money was actually traded.” The schemes of 2026 have updated the vocabulary, not the method.

In September the SEC charged two groups that sold access to AI trading bots and signals through WhatsApp; of one of them it wrote that “there were no AI trading bots and deposited funds never were used to earn returns for investors”[4].

On 29 October 2026 it sued a pair of companies that took more than $130 million for trading signals with a claimed accuracy of 99.6% and promised that “the monthly fund return rate is at least 60%”[45].

Sixty percent a month turns $1,000 into $281,000 in a year and into $22 million in three. A strategy that did this would not need investors.

The arithmetic of the earlier sections gives a better sense of scale: a model that is right 54% of the time on the next hour of Bitcoin, close to the best that published studies report, loses money at ordinary fees.

The losses are large. In 2025 Americans reported $7.2 billion lost to crypto investment fraud, the largest category of all in the FBI’s annual report, and more than $632 million in investment complaints where AI was mentioned[46].

Chainalysis found that scam operations that bought AI tools, from face swapping to text generation, earned 4.5 times as much as those that did not[47]. And in a survey by the UK’s Financial Conduct Authority, 44% of investors aged 18 to 40 believed that financial information generated by AI is regulated[48].

Strategies that do not need a forecast

Some of the steadiest money in crypto is made without predicting anything. A perpetual future is a contract that tracks a coin and never expires; to keep its price close to the coin’s, the side that is more crowded pays the other side every eight hours.

This payment is the funding rate. When more traders want to be long with leverage, which in crypto is most of the time, longs pay shorts.

Buying a coin and shorting the same amount of its perpetual cancels out the price risk and collects the funding. This is the cash-and-carry trade, and on Binance the rate defaults to 0.01% per eight hours, about 11% a year, when the two prices agree. What it actually paid looks like this:

Funding of Binance's BTC and ETH perpetual futures by month, scaled to a year, 2020 to October 2026. Yearly averages for BTC: 2020 17.2%, 2021 30.6%, 2022 4.2%, 2023 7.9%, 2024 11.9%, 2025 5.1%, 2026 3.1% Funding of Binance's BTC and ETH perpetual futures by month, scaled to a year, 2020 to October 2026. Yearly averages for BTC: 2020 17.2%, 2021 30.6%, 2022 4.2%, 2023 7.9%, 2024 11.9%, 2025 5.1%, 2026 3.1%

In 2021 the trade earned 30.6% a year on Bitcoin and 37.5% on Ether, from nothing but other traders’ appetite for leverage. Since then the Bitcoin trade has paid between 4 and 12% a year, and 3.1% in 2026 to the end of October, with a quarter of the payments going the other way.

A study at the BIS found crypto carry averaging more than 10% a year and sometimes reaching 60%[49]; the shrinking yield of recent years is what one would expect once a trade is easy to see and easy to automate.

Market-neutral is not risk-free. On the evening of 10 October 2025, after an announcement of new US tariffs on China, Bitcoin fell 12.6% on Binance in eighty minutes, Ether 14.1% and Solana 17.8%. More than $19 billion of leveraged positions were liquidated across exchanges.

Where insurance funds could not absorb the losses, exchanges closed profitable positions of other traders to cover them, a mechanism called auto-deleveraging, and among the positions closed were the short hedges of market makers[50].

On Hyperliquid alone it closed $2.1 billion of positions in twelve minutes[51]. A hedged trader whose hedge is closed by the exchange is no longer hedged.

Bitcoin falls 12.6% in eighty minuteslosses beyond the insurance fundsexchanges close the hedges of market makersa hedged trader is no longer hedged

The other risk is the exchange itself. A carry position lives on an exchange, and while I was writing this post the largest one stopped withdrawals after the hack mentioned above. A position that is neutral to the price of Bitcoin is not neutral to the solvency of the venue that holds it.

A realistic picture

Put together, the experiments draw a consistent picture, and it is less exciting than the advertising and more useful.

What machine learning does well in crypto is to forecast the things that have memory: volatility, volume, liquidity, the cost of trading.

These forecasts size positions, set quotes, route orders and keep a leveraged book alive through an hour like the evening of 10 October 2025. This is where the firms that make money from crypto markets put their models, and it rarely makes headlines.

Direction is a different matter. Short-horizon patterns exist, like the hourly reversal my model found, and other studies find them too: hit rates of 51 to 56% for Bitcoin from one minute to an hour ahead, with every model losing money once a round trip costs 30 bp[52]; on hourly Bitcoin bars from 2018 to 2026, models that earned 74 to 182% a year before costs lost 64 to 99% after a cost of 10 bp per trade[53].

Such an edge belongs to whoever trades most cheaply, which means market makers, not someone paying taker fees through an API.

Language models are good research assistants, programmers and readers of documents. As autonomous traders they have not shown skill in any test long and clean enough to tell, they remember the history they are tested on, and an agent with keys is a security problem before it is a trading strategy.

For an ML engineer moving into this field, the habits that matter are the ones markets punish for skipping: splitting by time with a gap, putting costs into the evaluation and not after it, counting how many variants were tried, thinking of a Sharpe ratio as an estimate with a standard error, and testing language models only on dates after their training.

For a trader, the same habits turn into questions to ask about any AI product. How many strategies were tried before this one was shown? How long is the live record, and does it include fees?

Were the backtest dates after the model’s training cutoff? Who holds the keys, and can they withdraw? And an offer of a fixed monthly return is not a trading product at all.

There is also a question for the market as a whole. Central banks have started to worry that many traders using the same few models will act alike; the Bank of England warned that more advanced AI trading strategies “could lead to firms taking increasingly correlated positions and acting in a similar way during a stress”[54].

In Alpha Arena, four of the six models, given the same prompt, moved together with hourly correlations between 0.52 and 0.83.

In a market where many agents read the same news with the same model, a crowded trade can form in minutes, and the hour that unwinds it is the kind that sits in the tails of the first chart of this post.

References

  1. Nof1 (2025). Exploring the limits of large language models as quant traders. Blog post, archived copy. ↩
  2. Azhang, J. (2025). Season 1 of Alpha Arena has officially ended. Post on X, 3 November 2025. ↩
  3. MEXC (2026). MEXC AI trading report: 85% use AI for crypto, yet only 6.5% would authorize it to trade most of their funds. Press release. ↩
  4. U.S. Securities and Exchange Commission (2026). SEC charges multiple entities in fraud schemes totaling at least $15 million that used WhatsApp and other platforms to lure investors. Press release 2026-95. ↩
  5. CoinDesk (2026). Binance uncovers wider breach in $6.3 billion exploit, targets Sunday withdrawal restart. 30 October 2026. ↩
  6. Corsi, F. (2009). A simple approximate long-memory model of realized volatility. Journal of Financial Econometrics, 7(2), 174–196. ↩
  7. Bailey, D. H., & López de Prado, M. (2014). The deflated Sharpe ratio: correcting for selection bias, backtest overfitting, and non-normality. Journal of Portfolio Management, 40(5), 94–107. ↩
  8. Harvey, C. R., Liu, Y., & Zhu, H. (2016). … and the cross-section of expected returns. Review of Financial Studies, 29(1), 5–68. ↩
  9. Lo, A. W. (2002). The statistics of Sharpe ratios. Financial Analysts Journal, 58(4), 36–52. ↩
  10. Fan, T., Yang, Y., Jiang, Y., Zhang, Y., Chen, Y., & Huang, C. (2025). AI-Trader: benchmarking autonomous agents in real-time financial markets. arXiv. ↩
  11. Chen, Y., et al. (2025). StockBench: can LLM agents trade stocks profitably in real-world markets?. arXiv. ↩
  12. Yu, H., Li, F., & You, J. (2025). LiveTradeBench: seeking real-world alpha with large language models. arXiv. ↩
  13. Barton, T. J., et al. (2026). What LLM trading agents actually do in production: a six-month, population-scale record from two fleets. arXiv. ↩
  14. Lopez-Lira, A., Tang, Y., & Zhu, M. (2025). The memorization problem: can we trust LLMs’ economic forecasts?. arXiv. ↩
  15. Li, X., Zeng, Y., Xing, X., Xu, J., & Xu, X. (2025). Profit Mirage: revisiting information leakage in LLM-based financial agents. arXiv. ↩
  16. Li, W. W., Kim, H., Cucuringu, M., & Ma, T. (2026). Can LLM-based financial investing strategies outperform the market in long run?. KDD 2026, Datasets and Benchmarks Track. ↩
  17. He, S., Lv, L., Manela, A., & Wu, J. (2025). Chronologically consistent large language models. arXiv. ↩
  18. Man Group (2025). What AI can (and can’t yet) do for alpha. Man Group insights. ↩
  19. Cointelegraph (2026). Over 95% of Coinbase’s code is written with AI: Rob Witoff. ↩
  20. Anthropic (2025). AI agents find smart contract exploits. Research post. ↩
  21. Fiszeder, P., Orzeszko, W., & Pietrzyk, R. (2026). News sentiment analysis using ChatGPT for Bitcoin price dynamics. Journal of Big Data, 13, 106. ↩
  22. Wintermute (2026). Machine learning researcher. Job advertisement on Built In. ↩
  23. Forecasting Research Institute (2026). AI models have likely reached parity with superforecasters on ForecastBench. 16 July 2026. ↩
  24. Saguillo, O., Ghafouri, V., Kiffer, L., & Suarez-Tangil, G. (2025). Unravelling the probabilistic forest: arbitrage in prediction markets. AFT 2025. ↩
  25. Shi, Y., Fu, Z., Chen, S., Zhao, B., Xu, W., Zhang, C., & Li, J. (2026). Kronos: a foundation model for the language of financial markets. AAAI 2026. ↩
  26. Ansari, A. F., et al. (2025). Chronos-2: from univariate to universal forecasting. arXiv. ↩
  27. Rahimikia, E., Ni, H., & Wang, W. (2025). Re(visiting) time series foundation models in finance. arXiv. ↩
  28. Brini, A. (2026). Forecasting realized volatility with time series foundation models: a comparison with econometric benchmarks. arXiv. ↩
  29. DL News (2024). Marc Andreessen distances himself from obscene memecoin GOAT after $300m surge. ↩
  30. CoinGecko (2025). AI agent market map: hype ends, technology continues. With Tiger Research. ↩
  31. CoinDesk (2026). AI agent token once worth $2.4 billion ends with founder calling it dead. 5 August 2026. ↩
  32. The Block (2024). Human player outwits Freysa AI agent in $47,000 crypto challenge. ↩
  33. The Block (2025). AI crypto bot AIXBT lost $100,000 worth of ETH after hacker gained unauthorized “dashboard access”. 18 March 2025. ↩
  34. CryptoSlate (2026). Grok’s crypto wallet was just exploited by a tweet sent in morse code without any private key compromise. ↩
  35. Patlan, A. S., Sheng, P., Hebbar, S. A., Mittal, P., & Viswanath, P. (2025). Real AI agents with fake memories: fatal context manipulation attacks on Web3 agents. arXiv. ↩
  36. PolySwarm (2026). The ClawHavoc campaign. ↩
  37. Knostic (2026). When “auto-signing” sends your wallet: a malicious MCP server on npm. ↩
  38. Liu, H., Shou, C., Wen, H., Chen, Y., Fang, R. J., & Feng, Y. (2026). Your agent is mine: measuring malicious intermediary attacks on the LLM supply chain. arXiv. ↩
  39. Rizvani, A., Apruzzese, G., & Laskov, P. (2026). Adversarial news and lost profits: manipulating headlines in LLM-driven algorithmic trading. IEEE SaTML 2026. ↩
  40. crypto.news (2026). Binance just gave AI bots a trading license. The safeguards are thinner than they look. ↩
  41. OKX (2026). Connecting agentic trading to OKX exchange with Agent Trade Kit. ↩
  42. Kraken (2026). Announcing the Kraken CLI. Kraken blog. ↩
  43. Commodity Futures Trading Commission (2024). Customer advisory: AI won’t turn trading bots into money machines. ↩
  44. Commodity Futures Trading Commission (2022). CFTC charges South African pool operator and CEO with $1.7 billion fraud involving bitcoin. Press release 8549-22. ↩
  45. U.S. Securities and Exchange Commission (2026). Complaint, SEC v. BG Wealth Sharing Ltd. and DSJ Exchange Pty Ltd.. U.S. District Court, Northern District of California, filed 29 October 2026. ↩
  46. FBI Internet Crime Complaint Center (2026). Internet crime report 2025. ↩
  47. Chainalysis (2026). 2026 crypto crime report: scams. ↩
  48. Financial Conduct Authority (2026). Young investors trust AI more than TV or celebrities. Press release. ↩
  49. Schmeling, M., Schrimpf, A., & Todorov, K. (2023). Crypto carry. BIS Working Papers, 1087. ↩
  50. CoinGlass (2026). CoinGlass 2025 crypto derivatives market annual report. ↩
  51. Chitra, T. (2025). Autodeleveraging: impossibilities and optimization. arXiv. ↩
  52. Jaquart, P., Dann, D., & Weinhardt, C. (2021). Short-term bitcoin market prediction via machine learning. Journal of Finance and Data Science, 7, 45–66. ↩
  53. Bysik, A., & Ślepaczuk, R. (2026). Machine learning-based bitcoin trading under transaction costs: evidence from walk-forward forecasting. arXiv. ↩
  54. Bank of England (2025). Financial stability in focus: artificial intelligence in the financial system. ↩
]]> Eva Koroleva Beyond attention with Mamba, state space models and hybrids 2026-10-19T19:00:00+00:00 2026-11-10T00:00:00+00:00 https://blog.korolevalab.com/blog/2026/transformer-architectures-10-mamba Transformer architectures: from BERT to diffusion · lesson 10 of 15 · all lessons · run on Kaggle · open in Colab · code on GitHub

Everything this course has built so far pays the same bill. Attention keeps every token it has read, in every layer, for as long as the conversation lasts (lesson 8), and compares every new token with all of them (lesson 9).

The bill grows with the text, and all the tricks of the last three lessons only make it grow more slowly.

A recurrent network doesn’t pay it. It reads a token, updates a state of fixed size, and forgets the token; the millionth token costs what the first one did. The trouble is that recurrent networks lost to Transformers for good reasons: they train one token at a time, and their state forgets what matters as easily as what doesn’t.

In December 2023 Albert Gu and Tri Dao, the same Tri Dao who wrote FlashAttention, published Mamba, a recurrent model that fixes both. It trains in parallel, like a Transformer, and its state decides, token by token, what to keep and what to let go.

The idea comes from an unexpected place, the state space models of control theory, the language of the Kalman filter that helped navigate Apollo to the Moon.

This lesson takes Mamba apart on the real weights of mamba-130m: the equation it starts from, how a continuous equation becomes a layer, what the model chooses to remember in a real sentence, the scan that makes a recurrence parallel, and generation with a state that never grows.

Then we look at what a fixed state costs, where Mamba ends up inside today’s hybrid models, and train two small ones from scratch to see what selectivity buys.

The same eight tokens read by attention and by Mamba. Attention keeps a key and a value for every token, so its cache grows by one slot per token and never shrinks. Mamba writes every token into one state of fixed size, 1,536 channels of 16 numbers per layer in mamba-130m, and the state is the same size after eight tokens or eight million The same eight tokens read by attention and by Mamba. Attention keeps a key and a value for every token, so its cache grows by one slot per token and never shrinks. Mamba writes every token into one state of fixed size, 1,536 channels of 16 numbers per layer in mamba-130m, and the state is the same size after eight tokens or eight million

What you’ll do

  1. Write a state space model from its differential equation and turn it into a layer with exact and approximate discretisation.
  2. Make it selective, as Mamba does, and write the whole Mamba block from scratch; load mamba-130m into it and match Hugging Face.
  3. Watch what the trained model writes into its state and how long each of its channels remembers.
  4. Write the scan twice, as a loop and as a parallel scan, and generate text with a state of fixed size, measured against a Transformer’s growing cache.

Setup

One small helper package, tfs, draws the course’s plots and token chips and knows where the notebook runs (Kaggle, Colab or your own machine). Everything the lesson teaches is written out in the notebook itself.

import importlib.util
import subprocess
import sys

TFS = "git+https://github.com/xmarva/transformers-from-scratch@main"
if importlib.util.find_spec("tfs") is None:
    subprocess.check_call([sys.executable, "-m", "pip", "install", "-q", TFS])

from tfs import env
from tfs.viz import table, token_chips, use_course_style

env.quiet()  # no progress bars or environment warnings in the outputs
SEED = 42
env.seed_everything(SEED)
COLORS = use_course_style()
WORK_DIR = env.work_dir()

import json
import math
import time

import matplotlib.pyplot as plt
from matplotlib.colors import LinearSegmentedColormap
import numpy as np
import pandas as pd
import torch
import torch.nn as nn
import torch.nn.functional as F
from huggingface_hub import hf_hub_download
from safetensors.torch import load_file
from transformers import AutoTokenizer

SMOKE = env.smoke()  # CI runs the lesson on a small slice; you run it in full
device = "cuda" if torch.cuda.is_available() else "cpu"
if device == "cuda":
    print("GPU:", torch.cuda.get_device_name(0))

GPU: Tesla T4

1. A state that doesn’t grow

A state space model describes a system by a vector of numbers, its state \(h(t) \in \mathbb{R}^N\), that changes over time under an input \(x(t)\) and produces an output \(y(t)\):

\[h'(t) = A\,h(t) + B\,x(t), \qquad y(t) = C\,h(t).\]

The first equation says how the state moves: \(A\) (\(N \times N\)) is how it would drift on its own, \(B\) (\(N \times 1\)) is how the input pushes it.

The second says what we read out of it. Rudolf Kálmán made this form the language of control theory in 1960, and it has run in autopilots, phone receivers and spacecraft since.

For a sequence model the useful case is a diagonal \(A\) with negative entries \(a_1, \dots, a_N\). The equations then split into \(N\) independent ones, and each is a leaky integrator:

\[h_n'(t) = a_n\,h_n(t) + b_n\,x(t).\]

Picture a bucket with a hole. The input pours water in; the hole lets a fraction \(\lvert a_n \rvert\) of whatever is inside leak out per unit of time. Without input, the level falls as \(e^{a_n t}\): a bucket with a small hole, \(a_n\) near zero, keeps its water for a long time; one with a big hole forgets in an instant. A state of \(N\) buckets with holes of different sizes keeps a record of the input at \(N\) different time scales, recent and distant, and \(C\) mixes the levels into an output.

Here are four such buckets, fed one splash of input at \(t = 1\) and a steady stream between \(t = 4\) and \(6\); the curves are the exact solution of each equation.

t = np.linspace(0, 10, 1001)
splash = (np.abs(t - 1) < 0.05) * 10.0  # a short burst whose total is 1
stream = ((t > 4) & (t < 6)) * 1.0
dt = t[1] - t[0]

fig, ax = plt.subplots(figsize=(7.4, 3.4))
colors = [COLORS["enc"], COLORS["att"], COLORS["data"], COLORS["dec"]]
for a, color in zip([-0.125, -0.25, -0.5, -1.0], colors):  # four holes, small to large
    level = np.zeros_like(t)
    for i in range(1, len(t)):  # exact step of h' = a h + x for constant x over dt
        x = splash[i - 1] + stream[i - 1]
        level[i] = level[i - 1] * math.exp(a * dt) + (math.exp(a * dt) - 1) / a * x
    ax.plot(t, level, color=color, linewidth=2, label=f"a = {a:g}")
ax.fill_between(t, 0, stream * 0.15, color=COLORS["line"], alpha=0.4, label="input")
ax.set_xlabel("time")
ax.set_ylabel("state")
ax.set_title("Four leaky buckets: the same input, four memories")
ax.legend(frameon=False, fontsize=9)
plt.show()

Four leaky buckets: the same input, four memories

The bucket with the smallest hole still remembers the splash when the stream arrives; the one with the largest has forgotten it within a unit of time and tracks only the present. That spread of memories is what a state space layer is made of.

2. From an equation to a layer

Text doesn’t arrive in continuous time; it arrives a token at a time. To turn the equation into a layer, assume the input stays constant for a short time \(\Delta\), the step, and solve the equation exactly over that time. This is the zero-order hold of signal processing, and for a diagonal \(A\) it gives, per channel,

\[h_t = \bar a\, h_{t-1} + \bar b\, x_t,\]

where

\[\bar a = e^{\Delta a}, \qquad \bar b = \frac{e^{\Delta a} - 1}{a}\, b.\]

\(\bar a\) is the fraction of the state that survives one step: between 0 and 1, because \(a < 0\). The state space model has become a linear recurrence, an RNN without a nonlinearity inside the loop. That missing nonlinearity is what will make it parallel in section 6.

Mamba’s code takes a shortcut for \(\bar b\). For small \(\Delta a\), \(e^{\Delta a} - 1 \approx \Delta a\), so \(\bar b \approx \Delta\, b\), the Euler step. Is that a big approximation? The ratio of the two is \((e^{\Delta a} - 1)/(\Delta a)\), which depends only on the product \(\Delta a\):

for product in [0.001, 0.01, 0.1, 0.5, 1.0, 2.0, 5.0]:
    exact_over_euler = (1 - math.exp(-product)) / product
    print(f"|Δa| = {product:<5}  exact b̄ / Euler b̄ = {exact_over_euler:.3f}")

|Δa| = 0.001  exact b̄ / Euler b̄ = 1.000
|Δa| = 0.01   exact b̄ / Euler b̄ = 0.995
|Δa| = 0.1    exact b̄ / Euler b̄ = 0.952
|Δa| = 0.5    exact b̄ / Euler b̄ = 0.787
|Δa| = 1.0    exact b̄ / Euler b̄ = 0.632
|Δa| = 2.0    exact b̄ / Euler b̄ = 0.432
|Δa| = 5.0    exact b̄ / Euler b̄ = 0.199

For a step that is small against the channel’s time scale, the two agree; for a large one the Euler step overestimates the input’s weight. We’ll measure in section 5 which regime the trained model lives in.

A model with fixed \(A\), \(B\), \(C\) and \(\Delta\) is linear and time-invariant: whatever the token, it is treated the same way.

S4, the 2021 model by Gu, Goel and Ré that started this line of work, was of this kind, and it could exploit that: a time-invariant recurrence is a convolution with a fixed kernel, and long convolutions are fast with the FFT.

S4 set records on the Long Range Arena, a benchmark of sequences up to 16k long, where the Transformers of its time did badly. On language it lagged behind. The reason is visible in the bucket picture: a fixed hole leaks the important word and the filler word at the same rate.

3. Selectivity

Mamba’s change is one sentence long: make \(\Delta\), \(B\) and \(C\) functions of the current token. Each token then decides, through its own \(\Delta_t\), how much of the state to keep and how much of itself to write:

\[h_t = e^{\Delta_t a}\, h_{t-1} + \Delta_t\, b_t\, x_t, \qquad y_t = c_t^{\top} h_t.\]

The step \(\Delta_t\) works as a gate. A large \(\Delta_t\) drives \(e^{\Delta_t a}\) toward 0 and multiplies the token’s input: the state is reset and the token is written in. A \(\Delta_t\) near zero keeps the state as it was and lets the token pass without trace.

Gu and Dao call this selectivity, and they show its point on a toy task, selective copying: remember the coloured tokens scattered among noise tokens and reproduce them at the end.

A time-invariant model can’t do it, because it can’t tell a coloured token from noise; a selective one can, because noise tokens get a \(\Delta\) near zero. Section 11 trains both kinds and checks.

Selectivity has a price: the recurrence is no longer a fixed convolution, so the FFT trick of S4 is gone. Mamba pays it with a scan, section 6, and with a GPU kernel written for it.

4. The Mamba block

The selective state space model is the core of a block that also contains a short convolution and a gate. One block, for a sequence of \(L\) vectors of width \(D\):

  1. in_proj widens each vector to \(2E\), with \(E = 2D\), and splits it into a stream \(u\) and a gate \(z\).
  2. conv1d mixes each channel of \(u\) with its three predecessors, a causal convolution of width 4, then SiLU. It gives the state space model a little local context to compute its \(\Delta\), \(B\), \(C\) from.
  3. x_proj maps each \(u_t\) to a low-rank seed of \(\Delta_t\) (rank \(R\)), and to \(b_t\) and \(c_t\), \(N\) numbers each, shared by all \(E\) channels.
  4. dt_proj expands the seed to one step per channel, and a softplus keeps the steps positive.
  5. The selective scan runs \(E\) independent channels, each with a state of \(N\) numbers, and adds a skip \(D \odot u\).
  6. The output is multiplied by \(\mathrm{SiLU}(z)\), the gate, and out_proj narrows it back to \(D\).

The Mamba block of mamba-130m with the shape of every tensor for L tokens. The input, L × 768, is widened by in_proj to L × 3,072 and split into a stream u and a gate z, L × 1,536 each. The stream goes through a causal convolution of width 4 and SiLU. x_proj turns it into a step seed of 48 numbers and B and C, 16 numbers each, per token; dt_proj and softplus give a step Δ for each of the 1,536 channels. The selective scan runs 1,536 channels with a state of 16 each, using A (1,536 × 16) and the skip D. Its output is multiplied by SiLU of the gate and out_proj narrows it back to L × 768 The Mamba block of mamba-130m with the shape of every tensor for L tokens. The input, L × 768, is widened by in_proj to L × 3,072 and split into a stream u and a gate z, L × 1,536 each. The stream goes through a causal convolution of width 4 and SiLU. x_proj turns it into a step seed of 48 numbers and B and C, 16 numbers each, per token; dt_proj and softplus give a step Δ for each of the 1,536 channels. The selective scan runs 1,536 channels with a state of 16 each, using A (1,536 × 16) and the skip D. Its output is multiplied by SiLU of the gate and out_proj narrows it back to L × 768

The matrix \(A\) is stored as A_log, \(a = -e^{A_{\log}}\), which keeps every entry negative whatever the training does. Mamba-130m has \(D = 768\), \(E = 1{,}536\), \(N = 16\), \(R = 48\) and 24 such blocks with no attention and no MLP between them: the block is the whole layer.

REPO = "state-spaces/mamba-130m-hf"
mamba_config = json.load(open(hf_hub_download(REPO, "config.json")))
for key in [
    "hidden_size",
    "intermediate_size",
    "state_size",
    "time_step_rank",
    "conv_kernel",
    "num_hidden_layers",
    "vocab_size",
]:
    print(f"{key:18} {mamba_config[key]}")

hidden_size        768
intermediate_size  1536
state_size         16
time_step_rank     48
conv_kernel        4
num_hidden_layers  24
vocab_size         50280

Two scans will be written in section 6; Mixer takes one as an argument. The loop is the definition: run the recurrence one step at a time, for all channels and all \(N\) state numbers at once.

def scan_loop(decay, write, c):
    # decay, write: (batch, channels, L, N); c: (batch, L, N) -> y: (batch, channels, L)
    h = torch.zeros(decay.shape[0], decay.shape[1], decay.shape[3], device=decay.device)
    ys = []
    for t in range(decay.shape[2]):
        h = decay[:, :, t] * h + write[:, :, t]
        ys.append((h * c[:, None, t]).sum(-1))
    return torch.stack(ys, dim=-1)


class RMSNorm(nn.Module):
    def __init__(self, size, eps):
        super().__init__()
        self.weight = nn.Parameter(torch.ones(size))
        self.eps = eps

    def forward(self, x):
        x32 = x.float()
        x32 = x32 * torch.rsqrt(x32.pow(2).mean(-1, keepdim=True) + self.eps)
        return self.weight * x32.to(x.dtype)

class Mixer(nn.Module):
    def __init__(self, c):
        super().__init__()
        d, e = c["hidden_size"], c["intermediate_size"]
        self.n, self.rank = c["state_size"], c["time_step_rank"]
        k = c["conv_kernel"]
        self.in_proj = nn.Linear(d, 2 * e, bias=c["use_bias"])
        self.conv1d = nn.Conv1d(e, e, k, groups=e, padding=k - 1, bias=c["use_conv_bias"])
        self.x_proj = nn.Linear(e, self.rank + 2 * self.n, bias=False)
        self.dt_proj = nn.Linear(self.rank, e, bias=True)
        self.A_log = nn.Parameter(torch.zeros(e, self.n))
        self.D = nn.Parameter(torch.ones(e))
        self.out_proj = nn.Linear(e, d, bias=c["use_bias"])
        self.scan = scan_loop
        self.keep = False  # section 5 reads the steps
        self.steps = None

    def forward(self, x):
        L = x.shape[1]
        u, gate = self.in_proj(x).transpose(1, 2).chunk(2, dim=1)  # (batch, E, L) each
        u = F.silu(self.conv1d(u)[..., :L])
        seed, b, c = self.x_proj(u.transpose(1, 2)).split(
            [self.rank, self.n, self.n], dim=-1
        )
        step = F.softplus(self.dt_proj(seed)).transpose(1, 2)  # (batch, E, L)
        if self.keep:
            self.steps = step.detach()
        a = -torch.exp(self.A_log.float())  # (E, N), every entry negative
        decay = torch.exp(step[..., None] * a[:, None, :])  # (batch, E, L, N)
        write = step[..., None] * b[:, None].float() * u[..., None].float()
        y = self.scan(decay, write, c.float()).to(u.dtype)
        y = y + u * self.D[None, :, None]
        y = y * F.silu(gate)
        return self.out_proj(y.transpose(1, 2))

class Block(nn.Module):
    def __init__(self, c):
        super().__init__()
        self.norm = RMSNorm(c["hidden_size"], c["layer_norm_epsilon"])
        self.mixer = Mixer(c)

    def forward(self, x):
        return x + self.mixer(self.norm(x))


class Mamba(nn.Module):
    def __init__(self, c):
        super().__init__()
        self.embeddings = nn.Embedding(c["vocab_size"], c["hidden_size"])
        self.layers = nn.ModuleList([Block(c) for _ in range(c["num_hidden_layers"])])
        self.norm_f = RMSNorm(c["hidden_size"], c["layer_norm_epsilon"])
        self.lm_head = nn.Linear(c["hidden_size"], c["vocab_size"], bias=False)

    def forward(self, ids):
        x = self.embeddings(ids)
        for layer in self.layers:
            x = layer(x)
        return self.lm_head(self.norm_f(x)).float()

Where do the parameters of a block go? Counted from the class:

block = Block(mamba_config)
rows = []
for name, p in block.mixer.named_parameters():
    rows.append(
        {
            "part": name,
            "shape": " × ".join(str(s) for s in p.shape),
            "parameters": p.numel(),
        }
    )
parts = pd.DataFrame(rows)
parts["share"] = parts.parameters / parts.parameters.sum()
table(parts, digits=3)
partshapeparametersshare
A_log1536 × 1624,5760.007
D15361,5360.000
in_proj.weight3072 × 7682,359,2960.626
conv1d.weight1536 × 1 × 46,1440.002
conv1d.bias15361,5360.000
x_proj.weight80 × 1536122,8800.033
dt_proj.weight1536 × 4873,7280.020
dt_proj.bias15361,5360.000
out_proj.weight768 × 15361,179,6480.313

The two projections, in and out, hold almost all of it, as the two matrices of an MLP do in a Transformer. The state space model itself, \(A\), \(D\), the convolution and the step’s projections, is a small fraction: Mamba’s work is done by a few thousand numbers per channel acting over time.

Now the weights. The library’s model runs in the same process here, because at 130 million parameters two copies fit anywhere; the smoke run uses a small random model of the same class.

from transformers import MambaConfig, MambaForCausalLM

if SMOKE:
    small = MambaConfig(
        vocab_size=mamba_config["vocab_size"],
        hidden_size=64,
        intermediate_size=128,
        state_size=8,
        time_step_rank=4,
        num_hidden_layers=2,
    )
    torch.manual_seed(0)
    reference = MambaForCausalLM(small).eval()
    mamba_config = small.to_dict()
    state = {k.removeprefix("backbone."): v for k, v in reference.state_dict().items()}
else:
    reference = MambaForCausalLM.from_pretrained(REPO, dtype=torch.float32).eval()
    state = load_file(hf_hub_download(REPO, "model.safetensors"))
    state = {k.removeprefix("backbone."): v for k, v in state.items()}
state.setdefault("lm_head.weight", state["embeddings.weight"])  # tied to the embeddings
mamba = Mamba(mamba_config)
mamba.load_state_dict(state, strict=True)
mamba.eval()

tok = AutoTokenizer.from_pretrained(REPO)


def tokens_of(ids):
    return [tok.decode([i]) for i in ids[0]]


sentence = (
    "In December 1911 Roald Amundsen reached the South Pole with four men and fifty-two "
    "dogs; Robert Falcon Scott arrived thirty-four days later and found a Norwegian tent."
)
ids = tok(sentence, return_tensors="pt").input_ids
with torch.no_grad():
    theirs = reference(ids).logits
    ours = mamba(ids)
print("tokens:", ids.shape[1])
print("largest logit:", f"{theirs.abs().max().item():.1f}")
print("max difference from Hugging Face:", f"{(ours - theirs).abs().max().item():.1e}")

tokens: 37
largest logit: 95.9
max difference from Hugging Face: 1.9e-04

Our loop and the library’s agree to the rounding of 32-bit floats accumulated over 24 layers. Every Mamba block is now ours to open.

One sentence through one block

The class hides the shapes, so here is the sentence going through the first block by hand, one stage at a time, with the shape of every tensor and a few of its numbers. Follow it with the block diagram above.

m = mamba.layers[0].mixer
with torch.no_grad():
    x = mamba.layers[0].norm(mamba.embeddings(ids))
    print("tokens in, normalised:          ", tuple(x.shape), "  (batch, L, D)")
    both = m.in_proj(x).transpose(1, 2)
    print("in_proj, channels first:        ", tuple(both.shape), "  (batch, 2E, L)")
    u, gate = both.chunk(2, dim=1)
    print("stream u and gate z:            ", tuple(u.shape), "each")
    L = ids.shape[1]
    u = F.silu(m.conv1d(u)[..., :L])
    print("after the convolution and SiLU: ", tuple(u.shape))
    seed, b, c = m.x_proj(u.transpose(1, 2)).split([m.rank, m.n, m.n], dim=-1)
    print(
        "x_proj → step seed, B, C:       ",
        tuple(seed.shape),
        tuple(b.shape),
        tuple(c.shape),
    )
    step = F.softplus(m.dt_proj(seed)).transpose(1, 2)
    print(
        "dt_proj, softplus → Δ:          ",
        tuple(step.shape),
        "  one step per channel and token",
    )
    a = -torch.exp(m.A_log.float())
    print("A, all negative:                ", tuple(a.shape))
    decay = torch.exp(step[..., None] * a[:, None, :])
    write = step[..., None] * b[:, None] * u[..., None]
    print("decay e^(Δa) and write ΔBu:     ", tuple(decay.shape), "  (batch, E, L, N)")
    y = scan_loop(decay, write, c)
    print("the scan, read out by C:        ", tuple(y.shape))
    y = (y + u * m.D[None, :, None]) * F.silu(gate)
    out = m.out_proj(y.transpose(1, 2))
    print(
        "skip D, gate, out_proj:         ",
        tuple(out.shape),
        "  back to the width of the model",
    )
    print("same as the block's own forward:", f"{(out - m(x)).abs().max():.1e}")

print()
word = 4
print(f"channel 0 at token {word} ({tokens_of(ids)[word]!r}):")
print("  step Δ:", f"{step[0, 0, word]:.4f}", "  share of the state kept:", end=" ")
print(" ".join(f"{v:.2f}" for v in decay[0, 0, word, :6].tolist()), "…")

tokens in, normalised:           (1, 37, 768)   (batch, L, D)
in_proj, channels first:         (1, 3072, 37)   (batch, 2E, L)
stream u and gate z:             (1, 1536, 37) each
after the convolution and SiLU:  (1, 1536, 37)
x_proj → step seed, B, C:        (1, 37, 48) (1, 37, 16) (1, 37, 16)
dt_proj, softplus → Δ:           (1, 1536, 37)   one step per channel and token
A, all negative:                 (1536, 16)
decay e^(Δa) and write ΔBu:      (1, 1536, 37, 16)   (batch, E, L, N)
the scan, read out by C:         (1, 1536, 37)
skip D, gate, out_proj:          (1, 37, 768)   back to the width of the model
same as the block's own forward: 0.0e+00

channel 0 at token 4 ('ald'):
  step Δ: 0.4137   share of the state kept: 0.00 0.23 0.96 0.25 0.20 0.00 …

5. What the model writes down

For each token and each of the 1,536 channels of a layer the model chooses a step \(\Delta_t\). A useful single number per token is how much of the state it replaces: \(1 - e^{\Delta_t a}\), averaged over channels and state numbers.

Zero means the token passed through without touching the state; one means the state was wiped and the token written instead.

for layer in mamba.layers:
    layer.mixer.keep = True
with torch.no_grad():
    mamba(ids)
tokens = tokens_of(ids)
replaced = []
for layer in mamba.layers:
    m = layer.mixer
    a = -torch.exp(m.A_log.detach().float())  # (E, N)
    kept = torch.exp(m.steps[0][:, :, None] * a[:, None, :])  # (E, L, N)
    replaced.append((1 - kept).mean(dim=(0, 2)))
    m.keep = False
replaced = torch.stack(replaced)  # (layers, L)
relative = (
    replaced / replaced.mean(dim=1, keepdim=True)
).numpy()  # 1 = the layer's average

show = list(range(0, len(mamba.layers), max(1, len(mamba.layers) // 6)))
fig, ax = plt.subplots(figsize=(9.6, 3.0))
shades = LinearSegmentedColormap.from_list("shades", [COLORS["card"], COLORS["att"]])
ax.imshow(relative[show], cmap=shades, aspect="auto", vmin=0.4, vmax=1.6)
ax.set_xticks(
    range(len(tokens)), [s.strip() or "␣" for s in tokens], rotation=70, fontsize=8
)
ax.set_yticks(range(len(show)), [f"layer {i}" for i in show])
ax.set_title("How much of the state each token replaces, against the layer's average")
plt.show()

middle = relative[len(relative) // 12 : len(relative) // 2 + 2].mean(axis=0)
order = np.argsort(-middle)
print("replace most: ", " · ".join(tokens[i].strip() for i in order[:8]))
print("replace least:", " · ".join(tokens[i].strip() for i in order[-8:]))

How much of the state each token replaces, against the layer

replace most:  In · December · the · Robert · Ro · thirty · a · four
replace least: days · en · mun · - · ds · - · four · two

Each layer is divided by its own average, because the early layers rewrite much more of their state than the late ones, and that difference would drown everything else. What remains is a pattern you can read.

The tokens that start something replace the most: the first word of the sentence, the first piece of a name, “Robert”, the month. The tokens that finish something replace the least: the remaining pieces of “Amundsen”, which the tokenizer cut into four, and the second half of “fifty-two” and “thirty-four”.

Once the state holds “Ro” and “ald”, the rest of the name is predictable and there is little new to write. Nobody told the model any of this; the steps were learned from predicting the next word.

The same steps give each channel a memory horizon: the number of tokens after which a channel keeps less than a third, \(e^{-1}\), of what it held. For a channel with entry \(a\) and a typical step \(\Delta\) that is \(1 / (\Delta \lvert a\rvert)\) tokens. Across all channels and state numbers of the whole model, over a page of text:

page = tok(
    "The expedition left the Bay of Whales on 20 October with four sledges. The weather "
    "was clear, the dogs were strong, and by the end of the first week they had covered "
    "more ground than Scott's party would in a month. Amundsen had planned the depots "
    "with care: every one was marked with flags across the line of march, so that a "
    "party returning in poor light could not miss it.",
    return_tensors="pt",
).input_ids
for layer in mamba.layers:
    layer.mixer.keep = True
with torch.no_grad():
    mamba(page)
horizons = []
for layer in mamba.layers:
    m = layer.mixer
    a = torch.exp(m.A_log.detach().float())  # |a|, (E, N)
    typical_step = m.steps[0].median(dim=-1).values  # (E,)
    horizons.append((1 / (typical_step[:, None] * a)).flatten())
    m.keep = False
horizons = torch.cat(horizons).clamp(max=1e6).numpy()

fig, ax = plt.subplots(figsize=(7.2, 3.2))
bins = np.logspace(-2, 6, 50)
ax.hist(horizons, bins=bins, color=COLORS["att"])
ax.set_xscale("log")
ax.set_xlabel("memory horizon, tokens")
ax.set_ylabel("channel × state pairs")
ax.set_title("How long the channels of mamba-130m remember")
plt.show()
for q in [0.1, 0.5, 0.9, 0.99]:
    print(f"{q:.0%} of horizons are under {np.quantile(horizons, q):,.2g} tokens")

How long the channels of mamba-130m remember

10% of horizons are under 0.03 tokens
50% of horizons are under 1.7 tokens
90% of horizons are under 4.2e+02 tokens
99% of horizons are under 2.7e+04 tokens

Most channels forget within a few tokens: they handle the local business of syntax. A long tail keeps its content for hundreds or thousands of tokens, and those are the ones that carry a topic or a name across a page.

The section 2 question has an answer too: for the short-memory channels \(\Delta \lvert a \rvert\) is large and the Euler step overstates the input, but those channels forget it in a token or two anyway; for the long-memory ones the product is small and the shortcut is exact for practical purposes.

6. The scan

The loop of section 4 is correct and slow: \(L\) steps, one after the other. A Transformer processes all positions of a sequence at once, and that is half of why it won. Mamba recovers the parallelism through a property of the recurrence \(h_t = \bar a_t h_{t-1} + w_t\): it is linear, so two consecutive steps compose into one step of the same form,

\[(\bar a_2, w_2) \circ (\bar a_1, w_1) = (\bar a_2 \bar a_1,\; \bar a_2 w_1 + w_2),\]

and this composition is associative. Associative operations can be accumulated over a sequence in \(\log_2 L\) rounds instead of \(L\) steps: in round \(k\), every position combines itself with the position \(2^k\) to its left.

After the last round every position holds the composition of everything before it. This is the Hillis–Steele scan from the parallel algorithms of the 1980s, and it needs only elementwise multiplications and additions.

def scan_parallel(decay, write, c):
    a, h = decay, write
    shift = 1
    while shift < a.shape[2]:
        a_left = F.pad(
            a[:, :, :-shift], (0, 0, shift, 0), value=1.0
        )  # identity beyond the start
        h_left = F.pad(h[:, :, :-shift], (0, 0, shift, 0), value=0.0)
        h = a * h_left + h
        a = a * a_left
        shift *= 2
    return (h * c[:, None]).sum(-1)


torch.manual_seed(SEED)
decay = torch.rand(1, 3, 37, 4)
write = torch.randn(1, 3, 37, 4)
c = torch.randn(1, 37, 4)
gap = (scan_parallel(decay, write, c) - scan_loop(decay, write, c)).abs().max()
print("parallel against loop:", f"{gap:.1e}")

parallel against loop: 9.5e-07

The Hillis–Steele scan on eight positions. In the first round every position combines with its left neighbour, in the second with the position two to its left, in the third with the position four to its left; after three rounds, the logarithm of eight, every position holds the combined effect of all the steps before it, and each round is one elementwise operation over the whole sequence The Hillis–Steele scan on eight positions. In the first round every position combines with its left neighbour, in the second with the position two to its left, in the third with the position four to its left; after three rounds, the logarithm of eight, every position holds the combined effect of all the steps before it, and each round is one elementwise operation over the whole sequence

The parallel scan does more work than the loop, \(L \log_2 L\) combinations instead of \(L\), and keeps a tensor of shape (batch, channels, \(L\), \(N\)) in memory, \(N\) times the size of the layer’s activations.

Written in plain PyTorch on a T4, the two come out about even: every step of the loop already runs over all 1,536 channels and 16 state numbers at once, so it is parallel across channels, and the scan’s extra work and memory traffic eat what it gains across time.

The scan pays off when it never leaves the chip, which is what Mamba’s own kernel does.

sizes = [256, 512] if SMOKE else [256, 1024, 4096]
timing = []
for L in sizes:
    torch.manual_seed(SEED)
    decay = torch.rand(1, 1536, L, 16, device=device)
    write = torch.randn(1, 1536, L, 16, device=device)
    c = torch.randn(1, L, 16, device=device)
    for name, scan in [("loop", scan_loop), ("parallel", scan_parallel)]:
        scan(decay, write, c)
        if device == "cuda":
            torch.cuda.synchronize()
        start = time.time()
        scan(decay, write, c)
        if device == "cuda":
            torch.cuda.synchronize()
        timing.append({"tokens": L, "scan": name, "ms": (time.time() - start) * 1000})
ms = pd.DataFrame(timing).pivot(index="tokens", columns="scan", values="ms")
table(ms.reset_index(), digits=1)
tokensloopparallel
25616.013.1
1,02463.462.8
4,096256.8303.8

Mamba’s own kernel goes further, and its design is FlashAttention’s lesson applied to a recurrence. The tensor of all states, (batch, channels, \(L\), \(N\)), is never written to the GPU’s main memory: the kernel loads \(\Delta\), \(A\), \(B\), \(C\) and \(x\) into the fast on-chip memory, runs the scan there and writes back only the outputs \(y\). For training it recomputes the states in the backward pass instead of storing them. With that kernel, the paper reports, the scan runs 20 to 40 times faster than a standard scan in PyTorch.

7. Generation with a state of fixed size

To generate, a Transformer keeps its KV cache; Mamba keeps, per layer, the state \(h\), \(E \times N\) numbers, and the last three inputs of the convolution. One step of generation reads the new token, updates both, and forgets the token. The step is the recurrence of section 3 for one position:

def mixer_step(m, x, state):
    window, h = (
        state  # the convolution's last inputs (batch, E, k), the state (batch, E, N)
    )
    u, gate = m.in_proj(x).chunk(2, dim=-1)  # (batch, E) each
    window = torch.cat([window[:, :, 1:], u[:, :, None]], dim=-1)
    u = F.silu((window * m.conv1d.weight[:, 0]).sum(-1) + m.conv1d.bias)
    seed, b, c = m.x_proj(u).split([m.rank, m.n, m.n], dim=-1)
    step = F.softplus(m.dt_proj(seed))  # (batch, E)
    a = -torch.exp(m.A_log.float())
    h = torch.exp(step[..., None] * a) * h + step[..., None] * b[:, None] * u[..., None]
    y = (h * c[:, None]).sum(-1) + u * m.D
    return m.out_proj(y * F.silu(gate)), (window, h)


def model_step(model, token, states):
    x = model.embeddings(token)
    new_states = []
    for layer, state in zip(model.layers, states):
        out, state = mixer_step(layer.mixer, layer.norm(x), state)
        x = x + out
        new_states.append(state)
    return model.lm_head(model.norm_f(x)).float(), new_states


def empty_states(model, batch=1):
    states = []
    for layer in model.layers:
        m = layer.mixer
        e, k = m.D.shape[0], m.conv1d.weight.shape[-1]
        dev = m.D.device
        states.append(
            (torch.zeros(batch, e, k, device=dev), torch.zeros(batch, e, m.n, device=dev))
        )
    return states


with torch.no_grad():
    full = mamba(ids)
    states = empty_states(mamba)
    stepped = []
    for t in range(ids.shape[1]):
        logits, states = model_step(mamba, ids[:, t], states)
        stepped.append(logits)
stepped = torch.stack(stepped, dim=1)
state_bytes = sum(w.numel() + h.numel() for w, h in states) * 2  # in 16-bit floats
print("step by step against the whole sequence:", f"{(stepped - full).abs().max():.1e}")
print(f"the whole state of the model: {state_bytes / 1e6:.2f} MB, at any length")

step by step against the whole sequence: 3.9e-04
the whole state of the model: 1.47 MB, at any length

Lesson 8 measured what a token costs in a Transformer’s cache. The same count for Pythia-160m, the Transformer that EleutherAI trained on the same 300 billion tokens of the Pile as mamba-130m, with the same tokenizer:

pythia_config = json.load(open(hf_hub_download("EleutherAI/pythia-160m", "config.json")))
per_token = (
    2 * pythia_config["num_hidden_layers"] * pythia_config["hidden_size"] * 2
)  # K and V, fp16
memory = []
for n in [2048, 32768, 131072, 1_000_000]:
    memory.append(
        {
            "tokens read": n,
            "Pythia-160m cache, MB": per_token * n / 1e6,
            "mamba-130m state, MB": state_bytes / 1e6,
        }
    )
table(pd.DataFrame(memory), digits=2)
tokens readPythia-160m cache, MBmamba-130m state, MB
2,04875.501.47
32,7681,207.961.47
131,0724,831.841.47
1,000,00036,864.001.47

The state is about as big as Pythia’s cache after a few dozen tokens, and it stays that size. This is the whole promise of the architecture, and it comes with a question that the next sections answer: what doesn’t fit into a state that never grows?

8. Mamba-2: the same model as attention

Six months after Mamba, Dao and Gu published Mamba-2 with a title that sounds like a provocation: Transformers are SSMs. The change to the model is small.

\(A\) becomes one number per head instead of one per channel and state number, the channels are grouped into heads of 64 like attention heads, \(B\) and \(C\) are computed in parallel with the input instead of after it, and the state grows from 16 numbers to 128.

The consequence is large. Unroll the recurrence with a scalar \(a\) per head and the output at position \(t\) is a weighted sum over all earlier positions:

\[y_t = \sum_{s \le t} \underbrace{c_t^{\top} b_s}_{\text{like } q_t^{\top} k_s}\; \underbrace{e^{\,a(\Delta_{s+1} + \dots + \Delta_t)}}_{\text{the mask}}\; \Delta_s\, x_s.\]

That is attention. \(c_t\) plays the query, \(b_s\) the key, \(\Delta_s x_s\) the value, and instead of a softmax the scores are multiplied by a mask that decays with distance at a rate the tokens choose.

The paper calls the equivalence state space duality: the same layer can be computed as a recurrence, cheap per token, or as a masked attention matrix, which is all matrix multiplication and therefore fast on a GPU.

The practical algorithm mixes the two: attention inside chunks of 64 to 256 tokens, the recurrence between chunks. Four steps, all of them einsum.

def segment_sums(a):
    # a (..., T) -> (..., T, T): out[t, s] = a[s+1] + ... + a[t] for s <= t,
    # and -inf above the diagonal
    T = a.shape[-1]
    below = torch.ones(T, T, dtype=torch.bool, device=a.device).tril(-1)
    x = a[..., None].expand(*a.shape, T).masked_fill(~below, 0)
    x = torch.cumsum(x, dim=-2)
    on_or_below = torch.ones(T, T, dtype=torch.bool, device=a.device).tril()
    return x.masked_fill(~on_or_below, float("-inf"))


def ssd(x, a, b, c, chunk=64):
    # x (batch, L, heads, P), already times Δ; a (batch, L, heads) = Δ·A;
    # b and c (batch, L, heads, N)
    batch, L, heads, p = x.shape
    pad = (chunk - L % chunk) % chunk
    x, b, c = [F.pad(t, (0, 0, 0, 0, 0, pad)) for t in (x, b, c)]
    a = F.pad(a, (0, 0, 0, pad))
    n = x.shape[1] // chunk
    x, b, c = [t.reshape(batch, n, chunk, heads, -1) for t in (x, b, c)]
    a = a.reshape(batch, n, chunk, heads).permute(
        0, 3, 1, 2
    )  # (batch, heads, chunks, chunk)
    a_cum = torch.cumsum(a, dim=-1)
    # 1. inside each chunk: attention with the decay as its mask
    mask = torch.exp(segment_sums(a))
    y_inside = torch.einsum("bclhn,bcshn,bhcls,bcshp->bclhp", c, b, mask, x)
    # 2. the state each chunk leaves behind
    to_end = torch.exp(a_cum[..., -1:] - a_cum)
    states = torch.einsum("bclhn,bhcl,bclhp->bchpn", b, to_end, x)
    # 3. the recurrence between chunks
    states = torch.cat([torch.zeros_like(states[:, :1]), states], dim=1)
    between = torch.exp(segment_sums(F.pad(a_cum[..., -1], (1, 0))))
    states = torch.einsum("bhzc,bchpn->bzhpn", between, states)[:, :-1]
    # 4. what each position reads from the state that entered its chunk
    y_entering = torch.einsum("bclhn,bchpn,bhcl->bclhp", c, states, torch.exp(a_cum))
    return (y_inside + y_entering).reshape(batch, n * chunk, heads, p)[:, :L]


def ssd_loop(x, a, b, c):
    h = torch.zeros(x.shape[0], x.shape[2], x.shape[3], b.shape[-1], device=x.device)
    ys = []
    for t in range(x.shape[1]):
        h = (
            torch.exp(a[:, t])[..., None, None] * h
            + x[:, t, :, :, None] * b[:, t, :, None, :]
        )
        ys.append((h * c[:, t, :, None, :]).sum(-1))
    return torch.stack(ys, dim=1)


torch.manual_seed(SEED)
x = torch.randn(1, 150, 3, 8)
a = -torch.rand(1, 150, 3) * 0.3
b = torch.randn(1, 150, 3, 16)
c = torch.randn(1, 150, 3, 16)
print(
    "chunks of 32 against the step-by-step loop:",
    f"{(ssd(x, a, b, c, 32) - ssd_loop(x, a, b, c)).abs().max():.1e}",
)

chunks of 32 against the step-by-step loop: 6.7e-06

The Mamba-2 block keeps the shape of section 4 with the changes above: one projection makes the gate, the stream, \(B\), \(C\) and the steps at once; the convolution runs over the stream together with \(B\) and \(C\); and the gate is applied inside an RMSNorm at the end. The checkpoint is the 130-million-parameter Mamba-2 released by the authors, in the Hugging Face format of a community conversion, AntonV/mamba2-130m-hf.

class GatedRMSNorm(nn.Module):
    def __init__(self, size, eps):
        super().__init__()
        self.weight = nn.Parameter(torch.ones(size))
        self.eps = eps

    def forward(self, x, gate):
        x32 = x.float() * F.silu(gate.float())
        x32 = x32 * torch.rsqrt(x32.pow(2).mean(-1, keepdim=True) + self.eps)
        return self.weight * x32.to(x.dtype)


class Mixer2(nn.Module):
    def __init__(self, c):
        super().__init__()
        d = c["hidden_size"]
        self.heads, self.p, self.n = c["num_heads"], c["head_dim"], c["state_size"]
        self.groups = c["n_groups"]
        self.e = self.heads * self.p
        conv = self.e + 2 * self.groups * self.n
        k = c["conv_kernel"]
        self.in_proj = nn.Linear(
            d, self.e + conv + self.heads, bias=c.get("use_bias", False)
        )
        self.conv1d = nn.Conv1d(conv, conv, k, groups=conv, padding=k - 1, bias=True)
        self.dt_bias = nn.Parameter(torch.zeros(self.heads))
        self.A_log = nn.Parameter(torch.zeros(self.heads))
        self.D = nn.Parameter(torch.ones(self.heads))
        self.norm = GatedRMSNorm(self.e, c["layer_norm_epsilon"])
        self.out_proj = nn.Linear(self.e, d, bias=c.get("use_bias", False))
        self.chunk = c.get("chunk_size", 64)
        self.keep = False  # the next cell reads the matrix of one head
        self.parts = None

    def select(self, dt, b, c):
        # the step, B and C of every token: what makes the layer selective (section 11)
        return F.softplus((dt + self.dt_bias).float()), b, c  # the step in 32 bits

    def forward(self, x):
        batch, L, _ = x.shape
        sizes = [self.e, self.e + 2 * self.groups * self.n, self.heads]
        gate, xbc, dt = self.in_proj(x).split(sizes, dim=-1)
        xbc = F.silu(self.conv1d(xbc.transpose(1, 2))[..., :L].transpose(1, 2))
        u, b, c = xbc.split([self.e, self.groups * self.n, self.groups * self.n], dim=-1)
        step, b, c = self.select(dt, b, c)
        a = -torch.exp(self.A_log.float())  # one number per head
        u = u.reshape(batch, L, self.heads, self.p)
        share = self.heads // self.groups
        b = b.reshape(batch, L, self.groups, self.n).repeat_interleave(share, dim=2)
        c = c.reshape(batch, L, self.groups, self.n).repeat_interleave(share, dim=2)
        if self.keep:
            self.parts = (u.detach(), step.detach(), a.detach(), b.detach(), c.detach())
        y = ssd(u.float() * step[..., None], step * a, b.float(), c.float(), self.chunk)
        y = y + u * self.D[:, None]
        y = self.norm(y.reshape(batch, L, self.e).to(x.dtype), gate)
        return self.out_proj(y)


class Mamba2(nn.Module):
    def __init__(self, c, mixer=Mixer2):
        super().__init__()
        self.embeddings = nn.Embedding(c["vocab_size"], c["hidden_size"])
        self.layers = nn.ModuleList()
        for _ in range(c["num_hidden_layers"]):
            block = nn.Module()
            block.norm = RMSNorm(c["hidden_size"], c["layer_norm_epsilon"])
            block.mixer = mixer(c)
            self.layers.append(block)
        self.norm_f = RMSNorm(c["hidden_size"], c["layer_norm_epsilon"])
        self.lm_head = nn.Linear(c["hidden_size"], c["vocab_size"], bias=False)

    def hidden(self, ids):
        x = self.embeddings(ids)
        for block in self.layers:
            x = x + block.mixer(block.norm(x))
        return self.norm_f(x)

    def forward(self, ids):
        return self.lm_head(self.hidden(ids)).float()

from transformers import Mamba2Config, Mamba2ForCausalLM

REPO2 = "AntonV/mamba2-130m-hf"
if SMOKE:
    small2 = Mamba2Config(
        vocab_size=50288,
        hidden_size=64,
        num_heads=4,
        head_dim=32,
        state_size=16,
        n_groups=1,
        num_hidden_layers=2,
        chunk_size=16,
    )
    torch.manual_seed(0)
    reference2 = Mamba2ForCausalLM(small2).eval()
    mamba2_config = small2.to_dict()
    state = {k.removeprefix("backbone."): v for k, v in reference2.state_dict().items()}
else:
    reference2 = Mamba2ForCausalLM.from_pretrained(REPO2, dtype=torch.float32).eval()
    mamba2_config = json.load(open(hf_hub_download(REPO2, "config.json")))
    state = load_file(hf_hub_download(REPO2, "model.safetensors"))
    state = {k.removeprefix("backbone."): v for k, v in state.items()}
state.setdefault("lm_head.weight", state["embeddings.weight"])
mamba2 = Mamba2(mamba2_config)
mamba2.load_state_dict(state, strict=True)
mamba2.eval()
with torch.no_grad():
    gap = (mamba2(ids) - reference2(ids).logits).abs().max()
print("Mamba-2, max difference from Hugging Face:", f"{gap:.1e}")
for key in ["num_heads", "head_dim", "state_size", "chunk_size"]:
    print(f"{key:12} {mamba2_config[key]}")
del reference2

Mamba-2, max difference from Hugging Face: 3.1e-04
num_heads    24
head_dim     64
state_size   128
chunk_size   256

Now the attention matrix that Mamba-2 never writes down. For one head of a middle layer, on the Amundsen sentence, the matrix \(M_{ts} = c_t^{\top} b_s\, e^{a(\Delta_{s+1} + \dots + \Delta_t)}\, \Delta_s\); next to it, the attention of a head of Pythia-160m on the same tokens. The check under the picture multiplies the matrix by the head’s values and compares with the scan.

layer2 = len(mamba2.layers) // 2
mixer2 = mamba2.layers[layer2].mixer
mixer2.keep = True
with torch.no_grad():
    mamba2(ids)
mixer2.keep = False
u, step, a, b, c = mixer2.parts
head = 0
L = ids.shape[1]
decay = torch.exp(segment_sums(step[0, :, head] * a[head]))  # (L, L)
M = (c[0, :, head] @ b[0, :, head].T) * decay * step[0, None, :, head]
y_matrix = M @ u[0, :, head]
y_scan = ssd(u * step[..., None], step * a, b, c, mixer2.chunk)[0, :, head]
print("matrix form against the chunked scan:", f"{(y_matrix - y_scan).abs().max():.1e}")

matrix form against the chunked scan: 7.6e-06
from transformers import AutoModelForCausalLM

PYTHIA = (
    "EleutherAI/pythia-14m" if SMOKE else "EleutherAI/pythia-160m"
)  # 14m: same tokenizer, tiny
pythia = AutoModelForCausalLM.from_pretrained(
    PYTHIA, dtype=torch.float32, attn_implementation="eager"
).eval()
with torch.no_grad():
    attention = pythia(
        ids, output_attentions=True
    ).attentions  # one (1, heads, L, L) per layer
pythia_head = attention[len(attention) // 2][0, 0].numpy()

labels = [s.strip() or "␣" for s in tokens_of(ids)]
fig, axes = plt.subplots(1, 2, figsize=(10.4, 5.0))
weights = M.abs().numpy()
weights = weights / weights.sum(
    axis=1, keepdims=True
)  # each row as a share, like attention
for ax, matrix, title in [
    (axes[0], weights, f"Mamba-2, layer {layer2}, head {head}: |M|, rows as shares"),
    (axes[1], pythia_head, f"Pythia-160m, layer {len(attention) // 2}, head 0: softmax"),
]:
    ax.imshow(np.sqrt(matrix), cmap=shades, vmin=0)
    ax.set_xticks(range(L), labels, rotation=90, fontsize=6)
    ax.set_yticks(range(L), labels, fontsize=6)
    ax.set_title(title, fontsize=9)
plt.show()

Output of the code above

Both matrices are lower-triangular: neither model looks ahead. The softmax head spends its attention on a few chosen tokens, often the first one, the attention sink of lesson 6.

Mamba-2’s matrix is a different texture: every row is a smooth decay into the past, interrupted where the steps \(\Delta\) are large, because a large step erases what came before. One is a lookup, the other a fading memory, and the next section measures what that difference means in practice.

9. Same data, different machinery

Pythia-160m and mamba-130m are as close to a controlled experiment as two public models get: the same 300 billion tokens of the Pile, the same tokenizer, the same training context of 2,048 tokens, and a similar number of parameters outside the embeddings. Only the layer differs. Two tests.

Reading past the training length. Lesson 9 watched a model with rotary positions fall off a cliff past its training length.

Pythia rotates a quarter of each head’s dimensions; Mamba has no positions at all, only the order in which the recurrence runs. Loss by position on the same PG-19 books as lesson 9:

import pandas as pd

pg19 = pd.read_parquet(
    "hf://datasets/emozilla/pg19-test/data/test-00000-of-00001-29a571947c0b5ccc.parquet"
)
LENGTH = 512 if SMOKE else 8192
N_BOOKS = 2 if SMOKE else 6
book_ids = []
for text in pg19.text:
    t = tok(text[:120_000], return_tensors="pt").input_ids
    if t.shape[1] > LENGTH:
        book_ids.append(t[:, : LENGTH + 1])
    if len(book_ids) == N_BOOKS:
        break

for layer in mamba.layers:
    layer.mixer.scan = scan_parallel
mamba.to(device)
pythia.to(device)


@torch.no_grad()
def position_losses(model_logits, ids):
    logits = model_logits(ids[:, :-1].to(device))
    return F.cross_entropy(logits[0], ids[0, 1:].to(device), reduction="none").cpu()


curves = {"mamba-130m": torch.zeros(LENGTH), "Pythia-160m": torch.zeros(LENGTH)}
for ids_book in book_ids:
    curves["mamba-130m"] += position_losses(mamba, ids_book) / len(book_ids)
    curves["Pythia-160m"] += position_losses(
        lambda x: pythia(x).logits.float(), ids_book
    ) / len(book_ids)

fig, ax = plt.subplots(figsize=(7.4, 3.4))
width = max(4, LENGTH // 64)
for name, color in [("Pythia-160m", COLORS["muted"]), ("mamba-130m", COLORS["att"])]:
    smooth = np.convolve(curves[name].numpy(), np.ones(width) / width, mode="valid")
    ax.plot(smooth, color=color, linewidth=2, label=name)
ax.axvline(2048, color=COLORS["line"], linestyle="--", linewidth=1)
ax.set_xlabel("position in the book")
ax.set_ylabel("loss, nats per token")
ax.set_title(f"{len(book_ids)} PG-19 books, both trained on 2,048 tokens")
ax.legend(frameon=False)
plt.show()
for name in curves:
    inside = curves[name][1024:2048].mean().item()
    beyond = curves[name][-2048:].mean().item()
    print(f"{name:12} loss at 1–2k: {inside:.2f}, at the last 2k: {beyond:.2f}")

Output of the code above

mamba-130m   loss at 1–2k: 3.05, at the last 2k: 3.16
Pythia-160m  loss at 1–2k: 3.22, at the last 2k: 7.65

Mamba reads eight times past its training length with its loss barely moving: there is no position for it to get wrong, only a state that keeps being updated. Pythia climbs the same wall as TinyLlama in lesson 9, from the same cause, rotary pairs asked for angles they never showed. On this test the fixed state wins outright.

Copying. The second test asks for exactly what a fixed state can’t promise: repeat a list of random words. A Transformer can look back at every word it was shown; Mamba has to have squeezed all of them into its state.

Jelassi and colleagues found in 2024 that even a 410-million-parameter Pythia beats a 2.8-billion-parameter Mamba at looking up a number in a long phone book, and the same gap shows on this simplest of tasks. Each prompt has one worked example and then a list of a given length; an answer counts when every word comes back in order.

import random

rng = random.Random(SEED)
words = sorted(
    {w.lower() for w in pg19.text.iloc[0].split() if w.isalpha() and 3 <= len(w) <= 7}
)


def copy_prompt(n_words):
    example = " ".join(rng.sample(words, 6))
    target = " ".join(rng.sample(words, n_words))
    text = f"Repeat the words exactly.\nWords: {example}\nRepeat: {example}\n\n"
    text += f"Words: {target}\nRepeat:"
    return text, target


@torch.no_grad()
def mamba_continue(model, ids, steps):
    states = empty_states(model)
    for t in range(ids.shape[1]):
        logits, states = model_step(model, ids[:, t].to(device), states)
    out = []
    for _ in range(steps):
        token = logits.argmax(-1)
        out.append(token.item())
        logits, states = model_step(model, token, states)
    return out


@torch.no_grad()
def pythia_continue(ids, steps):
    out = pythia.generate(
        ids.to(device), max_new_tokens=steps, do_sample=False, pad_token_id=tok.eos_token_id
    )
    return out[0, ids.shape[1] :].tolist()


copy_lengths = [5, 10] if SMOKE else [10, 25, 50, 100, 200]
COPY_TRIALS = 2 if SMOKE else 10
copied = {"mamba-130m": [], "Pythia-160m": []}
for n in copy_lengths:
    hits = {"mamba-130m": 0, "Pythia-160m": 0}
    for _ in range(COPY_TRIALS):
        text, target = copy_prompt(n)
        prompt = tok(text, return_tensors="pt").input_ids
        steps = len(tok(" " + target).input_ids) + 2
        said = {
            "mamba-130m": tok.decode(mamba_continue(mamba, prompt, steps)),
            "Pythia-160m": tok.decode(pythia_continue(prompt, steps)),
        }
        for name in hits:
            hits[name] += said[name].strip().startswith(target)
    for name in hits:
        copied[name].append(hits[name] / COPY_TRIALS)

fig, ax = plt.subplots(figsize=(7.2, 3.2))
for name, color in [("Pythia-160m", COLORS["muted"]), ("mamba-130m", COLORS["att"])]:
    ax.plot(copy_lengths, copied[name], color=color, marker="o", linewidth=2, label=name)
ax.set_xscale("log")
ax.set_xticks(copy_lengths, [str(n) for n in copy_lengths])
ax.set_ylim(-0.05, 1.05)
ax.set_xlabel("words to repeat")
ax.set_ylabel("lists repeated exactly")
ax.set_title(f"Repeating a list of random words, {COPY_TRIALS} lists per length")
ax.legend(frameon=False)
plt.show()

Output of the code above

On the shortest lists the two are close. From there Mamba falls behind fast and soon copies no list at all, while Pythia still gets some of the longer lists right, because every word it needs is still in its cache, one lookup away; Mamba’s state has to hold the whole list at once in a fixed number of numbers, and the longer the list, the less of each word survives.

Neither model is a good copyist at this size, and at the longest lists both fail, but the gap points the way Jelassi and colleagues found at scale. This is the trade at the heart of the architecture: a state that doesn’t grow is a state that can’t hold everything.

For most of language that is fine, because most of what we read doesn’t have to be recalled word for word. For the parts that do, a little attention goes a long way, which is the idea of the next section.

10. Hybrids

The answer the industry settled on is not “Mamba or attention” but a mixture: mostly recurrent layers for cheap long reading, and a few attention layers that can look anything up.

The Mamba-2 paper measured the ratio itself, on a 350-million-parameter model: replacing about one layer in ten with attention gave the best perplexity, better than either pure design. Here are the layouts of seven recent models, read from their configs. Each config describes its layout in its own way, so each needs a line of its own.

def layout(repo):
    c = read_hybrid_config(repo)
    n = c.get("num_hidden_layers") or c.get("n_layer")
    if "hybrid_override_pattern" in c:  # Nemotron-H: M Mamba, * attention, - an MLP layer
        names = {"M": "ssm", "*": "attention", "-": "mlp"}
        return [names[ch] for ch in c["hybrid_override_pattern"]]
    if "layer_types" in c:  # Granite 4.0
        return ["attention" if "attention" in t else "ssm" for t in c["layer_types"]]
    if "layers_block_type" in c:  # Zamba2: "hybrid" layers add a shared attention block
        return ["attention" if t == "hybrid" else "ssm" for t in c["layers_block_type"]]
    if "attn_layer_period" in c:  # Jamba
        period, offset = c["attn_layer_period"], c["attn_layer_offset"]
        return ["attention" if i % period == offset else "ssm" for i in range(n)]
    if "full_attention_interval" in c:  # Qwen3-Next: Gated DeltaNet, then attention
        k = c["full_attention_interval"]
        return ["attention" if (i + 1) % k == 0 else "linear" for i in range(n)]
    if "linear_attn_config" in c:  # Kimi Linear: its attention layers are counted from 1
        full = c["linear_attn_config"]["full_attn_layers"]
        return ["attention" if i + 1 in full else "linear" for i in range(n)]
    if "attn_layer_idx" in c:  # Mamba-2 with attention, from the Mamba-2 paper
        return ["attention" if i in c["attn_layer_idx"] else "ssm" for i in range(n)]
    return ["ssm"] * n


def read_hybrid_config(repo):
    with open(hf_hub_download(repo, "config.json")) as f:
        c = json.load(f)
    return c.get("text_config", c)


HYBRIDS = [
    ("state-spaces/mamba2attn-2.7b", "Mamba-2 + attention, 2024"),
    ("ai21labs/Jamba-v0.1", "Jamba, AI21, 2024"),
    ("Zyphra/Zamba2-1.2B", "Zamba2, Zyphra, 2024"),
    ("nvidia/Nemotron-H-8B-Base-8K", "Nemotron-H, NVIDIA, 2025"),
    ("ibm-granite/granite-4.0-h-tiny", "Granite 4.0 H, IBM, 2025"),
    ("Qwen/Qwen3-Next-80B-A3B-Instruct", "Qwen3-Next, 2025"),
    ("moonshotai/Kimi-Linear-48B-A3B-Instruct", "Kimi Linear, 2025"),
]
layouts = {name: layout(repo) for repo, name in HYBRIDS}
colors = {
    "attention": COLORS["dec"],
    "ssm": COLORS["att"],
    "linear": COLORS["enc"],
    "mlp": COLORS["line"],
}
fig, ax = plt.subplots(figsize=(10.4, 3.6))
for row, kinds in enumerate(layouts.values()):
    for col, kind in enumerate(kinds):
        ax.add_patch(plt.Rectangle((col, -row - 0.4), 0.9, 0.8, color=colors[kind]))
ax.set_xlim(0, max(len(k) for k in layouts.values()))
ax.set_ylim(-len(layouts) + 0.4, 0.6)
ax.set_yticks([-i for i in range(len(layouts))], list(layouts))
ax.set_xlabel("layer")
for kind, label in [
    ("attention", "attention"),
    ("ssm", "Mamba"),
    ("linear", "linear attention"),
    ("mlp", "MLP only"),
]:
    ax.plot([], [], color=colors[kind], linewidth=8, label=label)
ax.legend(frameon=False, ncol=4, loc="upper center", bbox_to_anchor=(0.5, 1.18))
plt.show()

rows = []
for name, kinds in layouts.items():
    mixing = [k for k in kinds if k != "mlp"]
    rows.append(
        {
            "model": name,
            "layers that mix tokens": len(mixing),
            "attention": mixing.count("attention"),
            "share of attention": mixing.count("attention") / len(mixing),
        }
    )
table(pd.DataFrame(rows), digits=2)

Output of the code above

modellayers that mix tokensattentionshare of attention
Mamba-2 + attention, 20246460.09
Jamba, AI21, 20243240.12
Zamba2, Zyphra, 20243860.16
Nemotron-H, NVIDIA, 20252840.14
Granite 4.0 H, IBM, 20254040.10
Qwen3-Next, 202548120.25
Kimi Linear, 20252770.26

The shares cluster where the Mamba-2 paper found its optimum: one attention layer in four to one in ten, spread evenly through the depth. Two other designs exist.

Falcon-H1 from TII runs attention and Mamba-2 heads in parallel inside every layer instead of alternating them, and Qwen3-Next and Kimi Linear replace Mamba with a different recurrence, the gated delta rule, which writes into its state more selectively; that family is the subject of a separate post on hybrid attention.

Not everyone stayed: MiniMax used linear attention in most layers of MiniMax-01 and went back to full attention in every layer for MiniMax-M2.

11. What selectivity buys

Section 3 claimed that a layer whose step, \(B\) and \(C\) depend on the token can do something a time-invariant one can’t: pick a few tokens out of a stream of noise and keep only those. Time to check, with selective copying, the toy task from the Mamba paper.

A sequence of 256 tokens is mostly noise; sixteen of its positions, chosen at random, hold coloured tokens, one of sixteen colours each; then come sixteen markers, and at the markers the model must say the coloured tokens back, in order. Guessing gets one in sixteen right.

Two models are trained from scratch, two Mamba-2 layers of width 64 each, identical except for one thing. In the first, the step, \(B\) and \(C\) come from each token, as in Mamba.

In the second, they are learned constants, the same for every token, which makes the state space part linear and time-invariant, the kind of layer S4 was. Mixer2 asks its select method for the three; the second model only replaces that method.

COLOURS, TO_COPY = 16, 16
MARKER = COLOURS + 1  # 0 is noise, 1 to 16 are the colours


def copying_batch(n, length):
    # TO_COPY coloured tokens at random places among `length` noise tokens, then markers
    where = torch.rand(n, length).argsort(dim=1)[:, :TO_COPY].sort(dim=1).values
    colours = torch.randint(1, COLOURS + 1, (n, TO_COPY))
    ids = torch.zeros(n, length + TO_COPY, dtype=torch.long)
    ids.scatter_(1, where, colours)
    ids[:, length:] = MARKER
    return ids, colours


class FixedMixer2(Mixer2):
    """The same layer with one step, B and C learned for all tokens, as in S4."""

    def __init__(self, c):
        super().__init__(c)
        self.b_fixed = nn.Parameter(torch.randn(self.groups * self.n) / self.n**0.5)
        self.c_fixed = nn.Parameter(torch.randn(self.groups * self.n) / self.n**0.5)

    def select(self, dt, b, c):
        step = F.softplus(self.dt_bias.float()).expand(dt.shape)
        return step, self.b_fixed.expand(b.shape), self.c_fixed.expand(c.shape)


def copier(mixer):
    c = dict(
        vocab_size=COLOURS + 2,
        hidden_size=64,
        num_heads=4,
        head_dim=32,
        state_size=32,
        n_groups=1,
        conv_kernel=4,
        num_hidden_layers=2,
        layer_norm_epsilon=1e-5,
        chunk_size=64,
    )
    torch.manual_seed(SEED)
    model = Mamba2(c, mixer)
    with torch.no_grad():  # Mamba's initialisation: several time scales, small steps
        for block in model.layers:
            m = block.mixer
            m.A_log.copy_(torch.log(torch.linspace(1, 16, m.heads)))
            step = torch.logspace(-3, -1, m.heads)
            m.dt_bias.copy_(step + torch.log(-torch.expm1(-step)))  # softplus inverse
    return model.to(device)


ids, colours = copying_batch(1, 64)
glyph = {0: "·", MARKER: "?"}
shown = " ".join(glyph.get(t, str(t)) for t in ids[0].tolist())
print("one sequence, 64 noise tokens shown as dots:")
print(shown[:94])
print(shown[94:])
print("the answer:", " ".join(map(str, colours[0].tolist())))

one sequence, 64 noise tokens shown as dots:
13 · · · · 1 · 10 14 · 1 · · · · · 10 · · 2 · · · · · 7 · · · · · · · · · · 16 10 · 11 13 9 · 
· 2 · · · · · · · · · 2 · · · · · · · 16 · ? ? ? ? ? ? ? ? ? ? ? ? ? ? ? ?
the answer: 13 1 10 14 1 10 2 7 16 10 11 13 9 2 2 16

The model answers the way a language model does, one position ahead: its output at the last noise token names the first colour, and its output at each marker the next one. Both models get the same budget: 6,000 steps of 64 sequences, a few minutes each on a T4.

COPY_LENGTH = 64 if SMOKE else 256
COPY_STEPS = 40 if SMOKE else 6000


def copying_accuracy(model, length, n=512):
    ids, colours = copying_batch(n, length)
    with torch.no_grad():
        guess = model(ids.to(device))[:, length - 1 : length + TO_COPY - 1].argmax(-1)
    return (guess.cpu() == colours).float().mean().item()


def train_copier(mixer):
    model = copier(mixer)
    opt = torch.optim.AdamW(model.parameters(), lr=3e-3, weight_decay=0.0)
    schedule = torch.optim.lr_scheduler.OneCycleLR(
        opt, max_lr=3e-3, total_steps=COPY_STEPS, pct_start=0.1
    )
    curve = []
    for step in range(1, COPY_STEPS + 1):
        ids, colours = copying_batch(64, COPY_LENGTH)
        logits = model(ids.to(device))[:, COPY_LENGTH - 1 : COPY_LENGTH + TO_COPY - 1]
        target = colours.to(device).reshape(-1)
        loss = F.cross_entropy(logits.reshape(-1, COLOURS + 2), target)
        opt.zero_grad()
        loss.backward()
        nn.utils.clip_grad_norm_(model.parameters(), 1.0)
        opt.step()
        schedule.step()
        if step % max(1, COPY_STEPS // 24) == 0:
            curve.append((step, copying_accuracy(model, COPY_LENGTH)))
    return model, curve


copiers, curves = {}, {}
for name, mixer in [("selective", Mixer2), ("fixed step, B and C", FixedMixer2)]:
    start = time.time()
    copiers[name], curves[name] = train_copier(mixer)
    print(f"{name:20} accuracy {curves[name][-1][1]:.3f}  in {time.time() - start:.0f} s")

fig, ax = plt.subplots(figsize=(6.4, 3.4))
for name, color in [("selective", COLORS["att"]), ("fixed step, B and C", COLORS["muted"])]:
    steps, accuracy = zip(*curves[name])
    ax.plot(steps, accuracy, color=color, marker="o", markersize=3, label=name)
ax.axhline(1 / COLOURS, color=COLORS["line"], linestyle=":", linewidth=1)
guess_y = 1 / COLOURS + 0.03
ax.text(COPY_STEPS, guess_y, "guessing", color=COLORS["muted"], ha="right", fontsize=9)
ax.set_ylim(0, 1.02)
ax.set_xlabel("training step")
ax.set_ylabel("coloured tokens copied")
ax.set_title(f"Selective copying, {COPY_LENGTH} tokens, two Mamba-2 layers")
ax.legend(frameon=False)
plt.show()

selective            accuracy 0.997  in 235 s
fixed step, B and C  accuracy 0.065  in 231 s

Output of the code above

The selective model learns the task within the budget, almost to the last token; the fixed one never leaves the guessing line. The paper trained a layer of this fixed kind far longer and it reached about half, but not more: whatever it does, it has to write every token into its state the same way, noise included.

Why does one work and the other not? The selective model’s own steps answer it. Here is one sequence, and for each position the step \(\Delta\) that the second layer chose, averaged over its four heads: a large step writes the token into the state and lets the old contents fade, a step near zero lets the token pass as if it weren’t there.

model = copiers["selective"]
mixer = model.layers[1].mixer
mixer.keep = True
ids, colours = copying_batch(1, COPY_LENGTH)
with torch.no_grad():
    model(ids.to(device))
mixer.keep = False
chosen = mixer.parts[1][0].mean(-1).cpu().numpy()  # (positions,), the step of every token
kind = ids[0].numpy()

fig, ax = plt.subplots(figsize=(11, 2.8))
for which, color, label in [
    (kind == 0, COLORS["line"], "noise"),
    ((kind > 0) & (kind <= COLOURS), COLORS["dec"], "coloured"),
    (kind == MARKER, COLORS["data"], "marker"),
]:
    where = np.flatnonzero(which)
    ax.bar(where, chosen[where], width=1.0, color=color, label=label)
fixed = copiers["fixed step, B and C"].layers[1].mixer
fixed_step = F.softplus(fixed.dt_bias).mean().item()
ax.axhline(fixed_step, color=COLORS["muted"], linestyle="--", linewidth=1)
note = "the fixed model's step, the same for every token"
ax.text(2, fixed_step * 1.15, note, color=COLORS["muted"], fontsize=9)
ax.set_xlabel("position in the sequence")
ax.set_ylabel("step Δ, layer 2")
ax.set_title("What the selective model writes down: only the coloured tokens")
ax.legend(frameon=False, ncol=3, loc="upper left")
plt.show()
for name, which in [("coloured", (kind > 0) & (kind <= COLOURS)), ("noise", kind == 0)]:
    print(f"mean step on {name} tokens: {chosen[which].mean():.4f}")

What the selective model writes down: only the coloured tokens

mean step on coloured tokens: 0.3641
mean step on noise tokens: 0.0297

Nobody told the model which tokens matter. It learned to give the coloured tokens large steps and the noise steps close to zero, so that the noise passes through without touching the state.

The exception is the very start, where noise gets large steps too: the state holds nothing yet, writing over it costs nothing, and training never had a reason to push those steps down.

The fixed model has one step for everything and writes the noise in as eagerly as the colours, and sixteen colours drown in hundreds of noise tokens.

A state that only opens for what matters doesn’t care how much noise surrounds it. The test below keeps the sixteen coloured tokens and grows the noise to two, four and eight times the length the models were trained on.

rows = []
for factor in [1, 2, 4, 8]:
    length = COPY_LENGTH * factor
    row = {"noise tokens": length}
    for name, model in copiers.items():
        row[name] = copying_accuracy(model, length, n=128)
    rows.append(row)
table(pd.DataFrame(rows), digits=3)
noise tokensselectivefixed step, B and C
2560.9970.068
5120.9990.060
1,0240.9920.061
2,0480.8410.068

Trained on 256 tokens, the selective model copies almost every colour at four times that length and most of them at eight: the noise it ignores costs it next to nothing. This is selectivity at its best, a few things worth keeping in a long stream.

Section 9 showed its limit: when every word of a long list has to be kept, a state of fixed size runs out, and that is the job the hybrids of section 10 hand to attention.

12. Where state space models are today

Pure Mamba at scale. Mistral released Codestral Mamba 7B in July 2024 and TII released Falcon Mamba 7B a month later, the first pure state space models of that size to compete with Transformers on standard benchmarks. Gu and Dao’s group went on to Mamba-3 in 2026, with a state that updates more expressively and decodes faster.

Hybrids as the default. Where long context and cheap serving matter, the shipped answer is a mixture: Jamba, Zamba2, Nemotron-H, Granite 4.0 H and Falcon-H1 keep a few attention layers for recall and use Mamba for the rest.

AI21 tried Mamba-2 for Jamba 1.5 and kept Mamba-1, which worked better next to attention; the newest version is not always the best partner.

Beyond text. Long sequences with simple alphabets are the natural home: DNA, audio and time series.

Mamba’s passport, the model that made state space models a language technology, and the family tree:

Mamba's passport. Published December 2023 by Albert Gu of Carnegie Mellon University and Tri Dao of Princeton University, code and weights under Apache 2.0. Type: a selective state space model, with no attention and no MLP blocks. Sizes: 130 million to 2.8 billion parameters. Training: 300 billion tokens of the Pile, context 2,048. State: per layer, 1,536 channels of 16 numbers in mamba-130m, plus a window of three inputs. Positions: none; order comes from the recurrence. Generation: the same memory and time for every token. Good at: long sequences at linear cost, fast generation, DNA and audio. Not for: recalling long spans exactly; pair it with attention Mamba's passport. Published December 2023 by Albert Gu of Carnegie Mellon University and Tri Dao of Princeton University, code and weights under Apache 2.0. Type: a selective state space model, with no attention and no MLP blocks. Sizes: 130 million to 2.8 billion parameters. Training: 300 billion tokens of the Pile, context 2,048. State: per layer, 1,536 channels of 16 numbers in mamba-130m, plus a window of three inputs. Positions: none; order comes from the recurrence. Generation: the same memory and time for every token. Good at: long sequences at linear cost, fast generation, DNA and audio. Not for: recalling long spans exactly; pair it with attention

The family of the Transformer as this course covers it, with lesson 10, Mamba and hybrids, highlighted and lessons 1 to 9 done The family of the Transformer as this course covers it, with lesson 10, Mamba and hybrids, highlighted and lessons 1 to 9 done

References

  1. Kalman, R. E. (1960). A new approach to linear filtering and prediction problems. Journal of Basic Engineering, 82(1), 35–45.
  2. Gu, A., Goel, K., & Ré, C. (2022). Efficiently modeling long sequences with structured state spaces. ICLR 2022.
  3. Gu, A., & Dao, T. (2023). Mamba: linear-time sequence modeling with selective state spaces. arXiv:2312.00752.
  4. Dao, T., & Gu, A. (2024). Transformers are SSMs: generalized models and efficient algorithms through structured state space duality. ICML 2024.
  5. Biderman, S., Schoelkopf, H., Anthony, Q., et al. (2023). Pythia: a suite for analyzing large language models across training and scaling. ICML 2023.
  6. Jelassi, S., Brandfonbrener, D., Kakade, S. M., & Malach, E. (2024). Repeat after me: Transformers are better than state space models at copying. arXiv:2402.01032.
  7. Arora, S., Eyuboglu, S., Timalsina, A., et al. (2023). Zoology: measuring and improving recall in efficient language models. arXiv:2312.04927.
  8. Lieber, O., Lenz, B., Bata, H., et al. (2024). Jamba: a hybrid Transformer–Mamba language model. arXiv:2403.19887.
  9. Hillis, W. D., & Steele, G. L. (1986). Data parallel algorithms. Communications of the ACM, 29(12), 1170–1183.
]]> Eva Koroleva Long context with RoPE scaling, YaRN and FlashAttention 2026-09-30T18:00:00+00:00 2026-11-10T00:00:00+00:00 https://blog.korolevalab.com/blog/2026/transformer-architectures-09-long Transformer architectures: from BERT to diffusion · lesson 9 of 15 · all lessons · run on Kaggle · open in Colab · code on GitHub

The context window is the one number on a model card that has grown faster than the parameter count. GPT-3 read 2,048 tokens and Llama 2 read 4,096; Llama 3.1 reads 131,072, and Gemini 1.5 and GPT-4.1 claim a million.

Almost none of these models spent their pretraining on windows that long. They were trained short and stretched afterwards, and the trick that made the stretching cheap came out of a Reddit forum.

In June 2023 a hobbyist who signs as kaiokendev divided every position number by four and got LLaMA to read 8,192 tokens after a little fine-tuning; Meta published the same idea as position interpolation within days, and a user called bloc97 posted a variant that needed no fine-tuning at all.

Two months later some of the same people wrote YaRN, and YaRN is what Qwen3, DeepSeek-V3 and gpt-oss switch on today to read past the length they were trained on.

Stretching the positions is half of the problem. The other half is that attention compares every token with every other one: a 64k-token prompt has four billion pairs per head, and writing their scores down would need more memory than a GPU has.

FlashAttention never writes them down. It computes attention one tile at a time and still gets the exact answer, thanks to a five-line trick for computing a softmax in one pass.

This lesson takes both halves apart. We watch TinyLlama, a model trained on 2,048 tokens, fall off a cliff at token 2,049, and find which of its rotary dimensions push it.

We write the standard fixes, a few lines each, check them against Hugging Face and race them on books. Then we build FlashAttention’s core from a softmax over five numbers, measure it on a T4, and finish with Qwen3 looking for a needle in a novel 64 thousand tokens long.

TinyLlama's 32 rotary pairs drawn as dials, fastest first, each showing how far the pair turned over the 2,048 positions of training. The fast pairs on the left turned hundreds of times, so every angle is familiar to the model. The slow pairs on the right never completed one turn: past position 2,048 they point in directions the model has never seen. Colors mark what YaRN does to each pair when it stretches the context eight times: the fast pairs are left alone, the slowest are stretched in full, and those in between are blended TinyLlama's 32 rotary pairs drawn as dials, fastest first, each showing how far the pair turned over the 2,048 positions of training. The fast pairs on the left turned hundreds of times, so every angle is familiar to the model. The slow pairs on the right never completed one turn: past position 2,048 they point in directions the model has never seen. Colors mark what YaRN does to each pair when it stretches the context eight times: the fast pairs are left alone, the slowest are stretched in full, and those in between are blended

What you’ll do

  1. Read from their configs how far eleven models read and how each one got there.
  2. Write one small decoder that runs both TinyLlama and Qwen3, and match Hugging Face.
  3. Watch TinyLlama’s loss explode past its training length, and find the rotary dimensions to blame.
  4. Write position interpolation, NTK scaling, dynamic NTK, YaRN and Llama 3.1’s scaling, check each against the transformers library, and race them on 16k-token books and a passkey test.
  5. Build FlashAttention’s forward pass from the online softmax, check it against plain attention, and measure memory and time on a T4.
  6. Hide a sentence in a novel and ask Qwen3-4B to find it, from 8k to 64k tokens, with and without YaRN.
  7. Give Qwen3 a whole novel, 77 thousand tokens, and ask it questions only a reader could answer.

What you need: a GPU. On Kaggle, Settings → Accelerator → GPU; on Colab, Runtime → Change runtime type → T4 GPU; the last two sections need two of them, GPU T4 x2 on Kaggle. The notebook downloads TinyLlama (4.4 GB), Qwen3-4B (8 GB), the hundred test books of PG-19 and one novel from Project Gutenberg.

Setup

One small helper package, tfs, draws the course’s plots and token chips and knows where the notebook runs (Kaggle, Colab or your own machine). Everything the lesson teaches is written out in the notebook itself.

import importlib.util
import subprocess
import sys

TFS = "git+https://github.com/xmarva/transformers-from-scratch@main"
if importlib.util.find_spec("tfs") is None:
    subprocess.check_call([sys.executable, "-m", "pip", "install", "-q", TFS])

from tfs import env
from tfs.viz import table, token_chips, use_course_style

env.quiet()  # no progress bars or environment warnings in the outputs
SEED = 42
env.seed_everything(SEED)
COLORS = use_course_style()
WORK_DIR = env.work_dir()

import json
import math
import random
import subprocess
import sys
import time
from pathlib import Path

import matplotlib.pyplot as plt
from matplotlib.colors import LinearSegmentedColormap
import numpy as np
import pandas as pd
import torch
import torch.nn as nn
import torch.nn.functional as F
from huggingface_hub import hf_hub_download
from safetensors.torch import load_file
from transformers import AutoTokenizer

SMOKE = env.smoke()  # CI runs the lesson on a small slice; you run it in full
device = "cuda" if torch.cuda.is_available() else "cpu"
HALF = torch.float16 if device == "cuda" else torch.float32  # 16-bit matmuls crawl on a CPU
assert device == "cuda" or SMOKE, "This lesson reads 64k-token prompts: use a GPU"
if device == "cuda":
    print("GPU:", torch.cuda.get_device_name(0))

GPU: Tesla T4

1. How far models read

A model’s config says how its positions are encoded and how far it is meant to read.

For a model with rotary positions, lesson 6’s RoPE, three things matter: the base of the rotations, the number of positions the model accepts, and, when the model was stretched after pretraining, how: the method, the factor and the length it was trained on. Here they are for eleven models.

LONG_MODELS = [  # repository, name; Llama and Gemma from Unsloth's public copies
    ("TinyLlama/TinyLlama-1.1B-intermediate-step-1431k-3T", "TinyLlama 1.1B"),
    ("unsloth/llama-2-7b", "Llama 2 7B"),
    ("unsloth/llama-3-8b", "Llama 3 8B"),
    ("unsloth/Meta-Llama-3.1-8B", "Llama 3.1 8B"),
    ("mistralai/Mistral-7B-v0.1", "Mistral 7B"),
    ("Qwen/Qwen2.5-7B-Instruct", "Qwen2.5 7B"),
    ("Qwen/Qwen3-8B", "Qwen3 8B"),
    ("unsloth/gemma-3-27b-pt", "Gemma 3 27B"),
    ("microsoft/Phi-3-mini-128k-instruct", "Phi-3 mini 128k"),
    ("deepseek-ai/DeepSeek-V3", "DeepSeek-V3"),
    ("openai/gpt-oss-20b", "gpt-oss 20B"),
]


def read_config(repo):
    with open(hf_hub_download(repo, "config.json")) as f:
        c = json.load(f)
    return c.get("text_config", c)


rows = []
for repo, name in LONG_MODELS:
    c = read_config(repo)
    stretch = c.get("rope_scaling") or {}
    trained_on = stretch.get("original_max_position_embeddings")
    rows.append(
        {
            "model": name,
            "RoPE base": c.get("rope_theta", 10000.0),  # absent: Llama's default
            "positions": c["max_position_embeddings"],
            "stretched by": stretch.get("rope_type") or stretch.get("type") or "—",
            "factor": stretch.get("factor", "—"),
            "trained on": trained_on or c.get("original_max_position_embeddings", "—"),
        }
    )
table(pd.DataFrame(rows), digits=0)
modelRoPE basepositionsstretched byfactortrained on
TinyLlama 1.1B10,0002,048———
Llama 2 7B10,0004,096———
Llama 3 8B500,0008,192———
Llama 3.1 8B500,000131,072llama388,192
Mistral 7B10,00032,768———
Qwen2.5 7B1,000,00032,768———
Qwen3 8B1,000,00040,960———
Gemma 3 27B1,000,000131,072linear8—
Phi-3 mini 128k10,000131,072longrope—4,096
DeepSeek-V310,000163,840yarn404,096
gpt-oss 20B150,000131,072yarn324,096

Two strategies show up. The first is to make the slowest rotation slow enough from the start: Llama went from a base of 10,000 to 500,000 between versions 2 and 3, Qwen to a million, and such models are then pretrained on long documents in a late stage of training.

The second is to train short and stretch: DeepSeek-V3 and gpt-oss were pretrained on 4,096-token windows and serve forty and thirty-two times more with YaRN; Llama 3.1 multiplies Llama 3’s 8,192 by eight with a scaling of Meta’s own; Phi-3 uses LongRoPE, a list of factors per dimension found by search.

Gemma 3 divides all positions by eight in its global layers, the simplest method of all. Most of the big ones do both.

A config’s number is a promise, not a measurement. Mistral 7B’s says 32,768, while the model was trained on 8,192 tokens with a sliding window of 4,096. Qwen3’s says 40,960, which its card explains as 32,768 tokens to read and 8,192 to write.

The cards of Qwen3-4B and the larger models add that YaRN with a factor of 4 reaches 131,072, and ask you to switch it on only when you need it, because a static stretch can hurt short texts. Section 8 tests that.

2. One decoder for two models

Two models carry the experiments. TinyLlama is Llama 2’s architecture at 1.1 billion parameters, trained by a team at the Singapore University of Technology and Design on three trillion tokens, all of them in windows of 2,048.

It is small enough to read 16k tokens on a T4 and old enough to have a short window: the ideal patient for the cliff and its cures.

Qwen3-4B is lesson 6’s model several sizes up, with a native window of 32k and a card that documents YaRN up to 131k; it is the patient for the needle and the novel.

The two differ in one detail: Qwen3 normalises each head’s queries and keys before rotating them, lesson 6’s QK-norm. One class covers both. RoPE first, written as in lesson 6 but with the frequencies passed in, because changing the frequencies is what this lesson is about.

Each pair of dimensions \(i\) of a head of size \(d\) turns by an angle \(\theta_i\) per position:

\[\theta_i = b^{-2i/d}, \qquad i = 0, 1, \dots, d/2 - 1,\]

where \(b\) is the base. At position \(m\) the pair is rotated by \(m\theta_i\). A scaling method replaces the \(\theta_i\) with new ones, and some also multiply the rotation by a constant, which scale carries.

def rope_frequencies(head_dim, base):
    pairs = torch.arange(0, head_dim, 2, device="cpu").float()  # not "meta" while loading
    return 1.0 / base ** (pairs / head_dim)


def rope_tables(freq, positions, scale=1.0):
    angles = positions[:, None].float() * freq[None, :].to(positions.device)
    angles = torch.cat([angles, angles], dim=-1)
    return angles.cos() * scale, angles.sin() * scale


def rotate(x, cos, sin):
    half = x.shape[-1] // 2
    turned = torch.cat([-x[..., half:], x[..., :half]], dim=-1)
    return x * cos + turned * sin

The cache is a block of memory allocated once for the longest text a run will read, with a pointer to where the next token goes. Lesson 8 grew its cache by concatenation, which copies everything stored so far on every step; at 64k tokens that copying would cost more than the attention itself.

class RMSNorm(nn.Module):
    def __init__(self, size, eps):
        super().__init__()
        self.weight = nn.Parameter(torch.ones(size))
        self.eps = eps

    def forward(self, x):
        x32 = x.float()
        x32 = x32 * torch.rsqrt(x32.pow(2).mean(-1, keepdim=True) + self.eps)
        return self.weight * x32.to(x.dtype)


class Cache:
    def __init__(self, size, kv_heads, head_dim, dtype, device):
        self.k = torch.empty(1, kv_heads, size, head_dim, dtype=dtype, device=device)
        self.v = torch.empty(1, kv_heads, size, head_dim, dtype=dtype, device=device)
        self.length = 0

    def add(self, k, v):
        end = self.length + k.shape[2]
        self.k[:, :, self.length : end] = k
        self.v[:, :, self.length : end] = v
        self.length = end
        return self.k[:, :, :end], self.v[:, :, :end]

Attention reads long texts in chunks of a few thousand tokens. A chunk sees every token before it and itself causally, which needs an explicit mask once earlier chunks sit in the cache.

The attention call itself is PyTorch’s scaled_dot_product_attention, which section 7 opens up. When one new token is generated, its queries are grouped by key-value head instead of copying the cache for every query head, which at 64k tokens would be most of the work.

class Attention(nn.Module):
    def __init__(self, c):
        super().__init__()
        self.heads = c["num_attention_heads"]
        self.kv_heads = c["num_key_value_heads"]
        self.head_dim = c.get("head_dim") or c["hidden_size"] // self.heads
        hidden = c["hidden_size"]
        bias = c.get("attention_bias", False)
        self.q_proj = nn.Linear(hidden, self.heads * self.head_dim, bias=bias)
        self.k_proj = nn.Linear(hidden, self.kv_heads * self.head_dim, bias=bias)
        self.v_proj = nn.Linear(hidden, self.kv_heads * self.head_dim, bias=bias)
        self.o_proj = nn.Linear(self.heads * self.head_dim, hidden, bias=False)
        self.qk_norm = c.get("model_type") == "qwen3"
        if self.qk_norm:
            self.q_norm = RMSNorm(self.head_dim, c["rms_norm_eps"])
            self.k_norm = RMSNorm(self.head_dim, c["rms_norm_eps"])
        self.keep = False  # section 6 looks at the last queries and all keys
        self.kept = None

    def forward(self, x, cos, sin, cache=None):
        b, t, _ = x.shape
        q = self.q_proj(x).view(b, t, self.heads, self.head_dim).transpose(1, 2)
        k = self.k_proj(x).view(b, t, self.kv_heads, self.head_dim).transpose(1, 2)
        v = self.v_proj(x).view(b, t, self.kv_heads, self.head_dim).transpose(1, 2)
        if self.qk_norm:
            q = self.q_norm(q)
            k = self.k_norm(k)
        q = rotate(q, cos, sin)
        k = rotate(k, cos, sin)
        if cache is not None:
            k, v = cache.add(k, v)
        if self.keep:
            self.kept = (q[:, :, -64:].float(), k.float())
        group = self.heads // self.kv_heads
        if t == 1:
            q = q.view(b, self.kv_heads, group, self.head_dim)
            scores = q @ k.transpose(-1, -2) / math.sqrt(self.head_dim)
            out = scores.float().softmax(-1).to(v.dtype) @ v
            out = out.reshape(b, self.heads, 1, self.head_dim)
        else:
            k = k.repeat_interleave(group, dim=1)
            v = v.repeat_interleave(group, dim=1)
            n = k.shape[2]
            if n == t:
                out = F.scaled_dot_product_attention(q, k, v, is_causal=True)
            else:
                query_pos = torch.arange(n - t, n, device=x.device)
                key_pos = torch.arange(n, device=x.device)
                mask = key_pos[None, :] <= query_pos[:, None]
                out = F.scaled_dot_product_attention(q, k, v, attn_mask=mask)
        return self.o_proj(out.transpose(1, 2).reshape(b, t, -1))

class MLP(nn.Module):
    def __init__(self, c):
        super().__init__()
        self.gate_proj = nn.Linear(c["hidden_size"], c["intermediate_size"], bias=False)
        self.up_proj = nn.Linear(c["hidden_size"], c["intermediate_size"], bias=False)
        self.down_proj = nn.Linear(c["intermediate_size"], c["hidden_size"], bias=False)

    def forward(self, x):
        return self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x))


class Block(nn.Module):
    def __init__(self, c):
        super().__init__()
        self.input_layernorm = RMSNorm(c["hidden_size"], c["rms_norm_eps"])
        self.self_attn = Attention(c)
        self.post_attention_layernorm = RMSNorm(c["hidden_size"], c["rms_norm_eps"])
        self.mlp = MLP(c)

    def forward(self, x, cos, sin, cache=None):
        x = x + self.self_attn(self.input_layernorm(x), cos, sin, cache)
        return x + self.mlp(self.post_attention_layernorm(x))

The model returns hidden states, not logits: a 64k-token prompt times Qwen3’s 151,936-word vocabulary is ten billion logits, and we only ever need a few of them.

class Decoder(nn.Module):
    def __init__(self, c):
        super().__init__()
        self.embed_tokens = nn.Embedding(c["vocab_size"], c["hidden_size"])
        self.layers = nn.ModuleList([Block(c) for _ in range(c["num_hidden_layers"])])
        self.norm = RMSNorm(c["hidden_size"], c["rms_norm_eps"])
        self.lm_head = nn.Linear(c["hidden_size"], c["vocab_size"], bias=False)
        self.head_dim = c.get("head_dim") or c["hidden_size"] // c["num_attention_heads"]
        self.base = c.get("rope_theta", 10000.0)  # absent from TinyLlama's config
        self.freq = rope_frequencies(self.head_dim, self.base)
        self.scale = 1.0

    def forward(self, ids, start=0, caches=None):
        x = self.embed_tokens(ids.to(self.embed_tokens.weight.device))
        positions = torch.arange(start, start + ids.shape[1], device=x.device)
        cos, sin = rope_tables(self.freq, positions, self.scale)
        for i, layer in enumerate(self.layers):
            here = layer.input_layernorm.weight.device  # a large model sits on two GPUs
            x = x.to(here)
            cache = caches[i] if caches is not None else None
            x = layer(x, cos.to(here, x.dtype), sin.to(here, x.dtype), cache)
        return self.norm(x.to(self.norm.weight.device))

    def logits(self, h):
        return self.lm_head(h.to(self.lm_head.weight.device)).float()

    def new_caches(self, size):
        caches = []
        for layer in self.layers:
            a = layer.self_attn
            w = layer.input_layernorm.weight
            caches.append(Cache(size, a.kv_heads, a.head_dim, w.dtype, w.device))
        return caches

Loading reads the safetensors files, one or several, and drops the model. prefix of the names. Qwen3-4B ties its output matrix to its embeddings and doesn’t store it twice.

With spread, a model too big for one T4 together with a long cache is cut in two: the first half of its layers on one GPU, the second half on the other, and the hidden state crosses over once per token on the way up.

def load_model(source, c, dtype, spread=False):
    if Path(source).is_dir():  # the smoke run's small models, saved to disk
        files = sorted(Path(source).glob("*.safetensors"))
    else:
        try:
            with open(hf_hub_download(source, "model.safetensors.index.json")) as f:
                names = sorted(set(json.load(f)["weight_map"].values()))
        except Exception:  # a single file has no index
            names = ["model.safetensors"]
        files = [hf_hub_download(source, name) for name in names]
    state = {}
    for path in files:
        for key, value in load_file(str(path)).items():
            state[key.removeprefix("model.")] = value.to(dtype)
    if "lm_head.weight" not in state:
        state["lm_head.weight"] = state["embed_tokens.weight"]
    with torch.device("meta"):
        model = Decoder(c)
    model.load_state_dict(state, strict=True, assign=True)
    model.eval()
    if not spread or torch.cuda.device_count() < 2:
        return model.to(device)
    half = len(model.layers) // 2  # half of the layers on each GPU
    model.embed_tokens.to("cuda:0")
    for i, layer in enumerate(model.layers):
        layer.to("cuda:0" if i < half else "cuda:1")
    model.norm.to("cuda:1")
    model.lm_head.to("cuda:1")
    return model

The check against Hugging Face runs the library’s model in a process of its own, as in lessons 7 and 8, so that two copies are never in memory together. The smoke run uses small models with random weights and the same code.

TINY = "TinyLlama/TinyLlama-1.1B-intermediate-step-1431k-3T"
QWEN = "Qwen/Qwen3-4B"
tiny_tok = AutoTokenizer.from_pretrained(TINY)
qwen_tok = AutoTokenizer.from_pretrained(QWEN)

if SMOKE:
    from transformers import LlamaConfig, LlamaForCausalLM, Qwen3Config, Qwen3ForCausalLM

    small = dict(hidden_size=64, intermediate_size=128, num_hidden_layers=2)
    small.update(num_attention_heads=4, num_key_value_heads=2, max_position_embeddings=64)
    torch.manual_seed(0)
    tiny_ref = LlamaForCausalLM(LlamaConfig(vocab_size=32000, **small))
    tiny_ref.save_pretrained(WORK_DIR / "tiny-small")
    qwen_config = Qwen3Config(
        vocab_size=151936, head_dim=16, tie_word_embeddings=True, **small
    )
    Qwen3ForCausalLM(qwen_config).save_pretrained(WORK_DIR / "qwen-small")
    TINY_SOURCE, QWEN_SOURCE = WORK_DIR / "tiny-small", WORK_DIR / "qwen-small"
    tiny_c = json.loads((TINY_SOURCE / "config.json").read_text())
    qwen_c = json.loads((QWEN_SOURCE / "config.json").read_text())
    for c in [tiny_c, qwen_c]:  # newer configs keep the base inside rope_parameters
        c["rope_theta"] = c.get("rope_parameters", {}).get("rope_theta", 10000.0)
else:
    TINY_SOURCE, QWEN_SOURCE = TINY, QWEN
    tiny_c = read_config(TINY)
    qwen_c = read_config(QWEN)

text = (
    "The lighthouse keeper on Skellig Michael logged every ship that passed in a green "
    "ledger, and in the winter of 1871 he logged none."
)
check_ids = tiny_tok(text, return_tensors="pt").input_ids
torch.save(check_ids, WORK_DIR / "check_ids.pt")
script = f"""
import torch
from transformers import AutoModelForCausalLM
device = "cuda" if torch.cuda.is_available() else "cpu"
dtype = torch.float16 if device == "cuda" else torch.float32
model = AutoModelForCausalLM.from_pretrained("{TINY_SOURCE}", dtype=dtype)
model = model.to(device).eval()
ids = torch.load("{WORK_DIR / "check_ids.pt"}").to(device)
with torch.no_grad():
    torch.save(model(ids).logits.float().cpu(), "{WORK_DIR / "theirs.pt"}")
"""
subprocess.run([sys.executable, "-c", script], check=True)
theirs = torch.load(WORK_DIR / "theirs.pt")

tiny = load_model(TINY_SOURCE, tiny_c, HALF)
with torch.no_grad():
    ours = tiny.logits(tiny(check_ids.to(device))).cpu()
print("tokens:", check_ids.shape[1])
print("largest logit:", f"{theirs.abs().max().item():.2f}")
print("max difference from Hugging Face:", f"{(ours - theirs).abs().max().item():.2e}")

tokens: 36
largest logit: 22.11
max difference from Hugging Face: 2.34e-02

The difference is the rounding of 16-bit floats over 22 layers; in 32-bit floats, on the smoke run’s random model, it is zero to the last bit. Qwen3 gets the same check in section 8.

3. The cliff at token 2,049

A language model’s loss on a token is how surprised it was by it. Averaged over many texts and plotted against the token’s position, it should fall slowly, because each new token has more context to lean on.

The texts are books from PG-19, Project Gutenberg books published before 1919, which DeepMind collected in 2019 to measure exactly this: how well a model uses long context. We take the first 16,384 tokens of eight of its test books.

pg19 = pd.read_parquet(
    "hf://datasets/emozilla/pg19-test/data/test-00000-of-00001-29a571947c0b5ccc.parquet"
)
TRAINED = tiny_c["max_position_embeddings"]
STRETCH = 8
LENGTH = STRETCH * TRAINED
N_BOOKS = 2 if SMOKE else 8
books = []
for book in pg19.itertuples():
    ids = tiny_tok(book.text[:250_000], return_tensors="pt").input_ids
    if ids.shape[1] > LENGTH:
        books.append(
            {"title": book.short_book_title, "ids": ids[:, : LENGTH + 1], "text": book.text}
        )
    if len(books) == N_BOOKS:
        break
for book in books:
    print(book["title"])

Reminiscences of Pioneer Days in St. Paul by Frank Moore
Dragon's blood by Henry Milner Rideout
Travels in Morocco Volume 2 by James Richardson
Impressions of Theophrastus Such by George Eliot
The S. W. F. Club by Caroline E. Jacobs
Frank Merriwell Down South by Burt L. Standish
Critical Miscellanies (Vol. 2 of 3) by John Morley
From Sand Hill to Pine by Bret Harte

token_losses reads a book in chunks through the cache and returns the loss of every token. stretch sets a model’s rotary frequencies by the name of a method; for now it knows only "none", the frequencies the model was trained with. Section 5 teaches it the rest.

@torch.no_grad()
def token_losses(model, ids, chunk=4096):
    n = ids.shape[1] - 1
    caches = model.new_caches(n)
    losses = []
    for a in range(0, n, chunk):
        b = min(a + chunk, n)
        h = model(ids[:, a:b].to(device), start=a, caches=caches)
        target = ids[0, a + 1 : b + 1].to(device)
        losses.append(F.cross_entropy(model.logits(h)[0], target, reduction="none").cpu())
    return torch.cat(losses)


def average_losses(model):
    total = torch.zeros(LENGTH)
    for book in books:
        total += token_losses(model, book["ids"])
    return total / len(books)


def smooth(curve, width=256):
    width = min(width, len(curve) // 8)
    return np.convolve(curve.numpy(), np.ones(width) / width, mode="valid")


STRETCHES = {}


def stretch(model, method, s=1.0):
    model.freq, model.scale = STRETCHES[method](model.head_dim, model.base, s)


def unchanged(head_dim, base, s):
    return rope_frequencies(head_dim, base), 1.0


STRETCHES["none"] = unchanged
stretch(tiny, "none")
curves = {"none": average_losses(tiny)}

fig, ax = plt.subplots(figsize=(7.4, 3.4))
ax.plot(smooth(curves["none"]), color=COLORS["dec"], linewidth=2)
ax.axvline(TRAINED, color=COLORS["line"], linestyle="--", linewidth=1)
top = ax.get_ylim()[1]
ax.text(TRAINED * 1.08, top * 0.92, f"trained length, {TRAINED:,}", color=COLORS["muted"])
ax.set_xlabel("position in the book")
ax.set_ylabel("loss, nats per token")
ax.set_title(f"TinyLlama on {len(books)} books, as it was trained")
plt.show()

inside = curves["none"][TRAINED // 2 : TRAINED].mean().item()
next_window = curves["none"][TRAINED : 2 * TRAINED].mean().item()
far = curves["none"][-TRAINED:].mean().item()
print(f"mean loss, second half of the trained window: {inside:.2f}")
print(f"mean loss, the next {TRAINED:,} tokens: {next_window:.2f}")
print(f"mean loss, the last {TRAINED:,} tokens: {far:.2f}")

Output of the code above

mean loss, second half of the trained window: 2.53
mean loss, the next 2,048 tokens: 9.30
mean loss, the last 2,048 tokens: 8.86

Up to the trained length the curve does what it should. Then it turns into a wall. Nothing in the weights changed at position 2,049; what changed is the angles the rotary pairs report.

4. Clocks that never finished a turn

Think of every rotary pair as a clock hand that advances by \(\theta_i\) per token. Its wavelength, the number of tokens for one full turn, is

\[\lambda_i = \frac{2\pi}{\theta_i} = 2\pi\, b^{2i/d},\]

and over a training window of \(L\) tokens the hand makes \(r_i = L / \lambda_i\) turns. Attention between tokens \(m\) and \(n\) sees, in pair \(i\), the angle \((m - n)\theta_i\). A pair that turned hundreds of times during training has shown the model every angle at every distance; at a distance of 3,000 it reports an angle the model knows well from a shorter distance. A pair that never finished a turn is different: in training it only ever reported angles between zero and \(L\theta_i\), and beyond \(L\) it reports angles the model has never seen. Those pairs carry the coarse sense of “how far back”, and they are the ones that lie.

turns = TRAINED * rope_frequencies(tiny.head_dim, tiny.base) / (2 * math.pi)
print("turns in training, fastest pair to slowest:")
print(" ".join(f"{t:.2f}" if t < 10 else f"{t:.0f}" for t in turns.tolist()))
slow = turns < 1
print(f"pairs that never finished a turn: {int(slow.sum())} of {len(turns)}")

turns in training, fastest pair to slowest:
326 244 183 137 103 77 58 43 33 24 18 14 10 7.73 5.80 4.35 3.26 2.44 1.83 1.37 1.03 0.77 0.58 0.43 0.33 0.24 0.18 0.14 0.10 0.08 0.06 0.04
pairs that never finished a turn: 11 of 32

One pair at a time makes this concrete. Take three of TinyLlama’s pairs, a fast one, a middle one and a slow one, and ask what angle each reports between two tokens 16,000 positions apart. The shaded arc is every angle the pair showed during training, at distances from 0 to 2,048; the red hand is the angle at 16,000.

def draw_clocks(freq, stretched=None):
    distance = min(16000, LENGTH - 1)
    pairs = [2, len(freq) // 2, len(freq) - 4]
    fig, axes = plt.subplots(1, 3, figsize=(8.4, 3.0), subplot_kw={"projection": "polar"})
    for ax, pair in zip(axes, pairs):
        seen = TRAINED * freq[pair].item()  # the widest angle training showed, in radians
        arc = np.linspace(0, min(seen, 2 * math.pi), 200)
        ax.fill_between(arc, 0, 1, color=COLORS["line"], alpha=0.5)
        hands = [(freq, COLORS["dec"])]
        if stretched is not None:
            hands.append((stretched, COLORS["enc"]))
        for f, color in hands:
            angle = (distance * f[pair].item()) % (2 * math.pi)
            ax.plot([angle, angle], [0, 1], color=color, linewidth=2.5)
        ax.set_yticks([])
        ax.set_xticks([])
        turns_seen = seen / (2 * math.pi)
        ax.set_title(f"pair {pair}: {turns_seen:.2f} turns in training", fontsize=10)
    plt.show()
    for pair in pairs:
        seen = math.degrees(TRAINED * freq[pair].item())
        asked = math.degrees(distance * freq[pair].item())
        line = f"pair {pair:2}: training showed up to {min(seen, 360):6.1f}°;"
        line += f" at {distance:,} it asks for {asked:9.1f}°"
        if stretched is not None:
            after = math.degrees(distance * stretched[pair].item())
            line += f", after the stretch {after:9.1f}°"
        print(line)


draw_clocks(rope_frequencies(tiny.head_dim, tiny.base))

Output of the code above

pair  2: training showed up to  360.0°; at 16,000 it asks for  515516.6°
pair 16: training showed up to  360.0°; at 16,000 it asks for    9167.3°
pair 28: training showed up to   37.1°; at 16,000 it asks for     289.9°

The fast pair asks for an angle of many turns, which lands somewhere on a circle it has covered completely: familiar. The slow pair asks for an angle far beyond the sliver it ever showed: a direction the model has never seen.

That is the picture at the top of the lesson. If the story is right, slowing down only the pairs that never finished a turn, so that 16k tokens turn them no further than 2,048 used to, should take the wall down, and slowing down only the fast ones should not. Both experiments take a few lines.

def slow_pairs_only(head_dim, base, s):
    freq = rope_frequencies(head_dim, base)
    return torch.where(slow, freq / s, freq), 1.0


def fast_pairs_only(head_dim, base, s):
    freq = rope_frequencies(head_dim, base)
    return torch.where(slow, freq, freq / s), 1.0


STRETCHES["slow pairs only"] = slow_pairs_only
STRETCHES["fast pairs only"] = fast_pairs_only
for method in ["slow pairs only", "fast pairs only"]:
    stretch(tiny, method, STRETCH)
    curves[method] = average_losses(tiny)


def bucket_table(names):
    edges = [0, TRAINED // 2, TRAINED, 2 * TRAINED, LENGTH // 2, LENGTH]
    rows = []
    for name in names:
        row = {"stretch": name}
        for a, b in zip(edges[:-1], edges[1:]):
            row[f"{a:,}–{b:,}"] = curves[name][a:b].mean().item()
        rows.append(row)
    return table(pd.DataFrame(rows), digits=2)


bucket_table(["none", "slow pairs only", "fast pairs only"])
stretch0–1,0241,024–2,0482,048–4,0964,096–8,1928,192–16,384
none2.252.539.308.898.92
slow pairs only2.705.435.575.825.90
fast pairs only5.365.527.637.908.15

The table is the loss averaged over stretches of positions, the first two columns inside the trained window and the other three beyond it.

The fast pairs are not the problem: slowing them down leaves the wall standing and spoils the trained window as well. Slowing down the slow ones lowers the wall a long way, so that is where it comes from.

But it is a clumsy cure, and it costs a lot even inside the window, where nothing was wrong. It cuts the spectrum in two: the last pair that finished a turn keeps its speed, and its neighbour suddenly turns eight times slower.

The model reads neighbouring pairs together, and a seam between them hurts. The methods of the next section are, one way or another, ways to slow the slow pairs down without a seam.

5. Five ways to stretch the clocks

Position interpolation (Chen, Wong, Chen and Tian at Meta; kaiokendev at the same time). To read \(s\) times further, divide every position by \(s\):

\[m\theta_i \;\to\; \frac{m}{s}\theta_i \quad\Longleftrightarrow\quad \theta_i \;\to\; \frac{\theta_i}{s}.\]

Every angle at position 16,384 is now an angle the model saw at 2,048, so nothing is out of range. The fast pairs pay for it. Two neighbouring tokens used to differ by \(\theta_0 = 1\) radian in the fastest pair and now differ by \(1/8\), and the model’s sense of word order, which lives in those pairs, gets blurry. The paper’s answer is a short fine-tuning, after which the model has learned the new spacing.

NTK-aware scaling (bloc97 on Reddit). Leave the fast pairs alone and stretch the slow ones, by changing the base instead of the positions:

\[b' = b \cdot s^{\,d/(d-2)},\]

which gives every pair a new frequency

\[\theta'_i = \theta_i \cdot s^{-2i/(d-2)}.\]

The fastest pair, \(i = 0\), keeps its frequency; the slowest, \(i = d/2 - 1\), is divided by exactly \(s\); the pairs in between are divided by a power of \(s\) that grows with \(i\). The odd exponent \(d/(d-2)\) is what makes the slowest pair land exactly where position interpolation would put it.

The name points to the theory of neural tangent kernels and to the Fourier-features work of Tancik and colleagues, which bloc97 cited: networks learn high frequencies poorly unless they are handed them, so the fast pairs shouldn’t be squeezed.

Dynamic NTK (emozilla on Reddit). Don’t stretch until you have to: recompute the base with \(s = n / L\) as the text grows past \(L\), so that short texts run unchanged. It shares a catch with every scheme that changes while you read: the keys already in the cache were rotated with the old base.

def interpolate(head_dim, base, s):
    return rope_frequencies(head_dim, base) / s, 1.0


def ntk(head_dim, base, s):
    return rope_frequencies(head_dim, base * s ** (head_dim / (head_dim - 2))), 1.0


STRETCHES["PI"] = interpolate
STRETCHES["NTK"] = ntk

YaRN (Peng, Quesnelle, Fan and Shippole, 2023) does two things. The first is NTK-by-parts, the experiment of section 4 turned into a rule. Measure each pair by the turns \(r_i = L/\lambda_i\) it made in training.

Pairs that made more than \(\beta = 32\) turns are left alone; pairs that made fewer than \(\alpha = 1\) are interpolated in full; the ones in between are blended:

\[\theta'_i = \big(1 - \gamma_i\big)\,\frac{\theta_i}{s} + \gamma_i\,\theta_i,\]

where \(\gamma_i\) goes from 1 for the fast pairs to 0 for the slow ones. The second is a temperature: all attention logits are multiplied by \(1/t\), with

\[\sqrt{1/t} = 0.1 \ln s + 1.\]

A longer text gives every query more keys to spread its attention over, and its softmax gets flatter; the temperature sharpens it back. The formula is a fit the authors found empirically.

Implementations multiply both the queries and the keys by \(\sqrt{1/t}\), which is why the formula is written for the square root. Section 6 measures the flattening.

Now a detail on which the paper and the code disagree. The paper draws \(\gamma\) as a straight line in the number of turns, \(\gamma_i = (r_i - \alpha)/(\beta - \alpha)\).

The authors’ code, which Hugging Face and vLLM follow, draws the straight line in the index of the pair: it finds the pair that makes \(\beta\) turns and the pair that makes \(\alpha\), rounds the first down and the second up to whole pairs, and blends linearly between the two indices.

Turns fall exponentially with the index, so these are different curves. We write the code’s version, because that is what every model with YaRN in its config runs, and keep the paper’s for comparison.

def yarn(head_dim, base, s, trained=None, fast=32, slow=1):
    trained = trained or TRAINED
    freq = rope_frequencies(head_dim, base)

    def pair_making(turns):  # the pair index that makes this many turns in training
        return head_dim * math.log(trained / (turns * 2 * math.pi)) / (2 * math.log(base))

    first = max(math.floor(pair_making(fast)), 0)
    last = min(math.ceil(pair_making(slow)), head_dim - 1)
    pairs = torch.arange(head_dim // 2).float()
    gamma = 1 - ((pairs - first) / (last - first)).clamp(0, 1)  # 1: keep, 0: interpolate
    mixed = (1 - gamma) * freq / s + gamma * freq
    return mixed, 0.1 * math.log(s) + 1


def yarn_as_written(head_dim, base, s, trained=None, alpha=1, beta=32):
    trained = trained or TRAINED
    freq = rope_frequencies(head_dim, base)
    turns = trained * freq / (2 * math.pi)
    gamma = ((turns - alpha) / (beta - alpha)).clamp(0, 1)
    return (1 - gamma) * freq / s + gamma * freq, 0.1 * math.log(s) + 1


def yarn_cold(head_dim, base, s):
    freq, _ = yarn(head_dim, base, s)
    return freq, 1.0


STRETCHES["YaRN"] = yarn
STRETCHES["YaRN, no temperature"] = yarn_cold

Back to the three clocks, now with YaRN’s hand in teal: the fast pair keeps its angle, the slow pair’s hand returns into the arc the model knows.

yarn_freq, _ = yarn(tiny.head_dim, tiny.base, STRETCH)
draw_clocks(rope_frequencies(tiny.head_dim, tiny.base), yarn_freq)

Output of the code above

pair  2: training showed up to  360.0°; at 16,000 it asks for  515516.6°, after the stretch  515516.6°
pair 16: training showed up to  360.0°; at 16,000 it asks for    9167.3°, after the stretch    4231.1°
pair 28: training showed up to   37.1°; at 16,000 it asks for     289.9°, after the stretch      36.2°

Llama 3.1 stretches Llama 3 eight times with a rule of Meta’s own, which the Llama 3 paper doesn’t describe: it lives only in Meta’s code, with its constants commented “values obtained from grid search”.

Pairs whose wavelength is shorter than a quarter of the training window are kept, pairs whose wavelength is longer than the window are divided by the factor, and the ones in between are blended linearly in \(L/\lambda\), that is, in the number of turns.

That is the paper’s version of YaRN’s ramp with \(\alpha = 1\) and \(\beta = 4\), without the temperature. The code below follows Meta’s, and the check after it shows the coincidence is exact.

def llama3(head_dim, base, s, trained=None, low=1, high=4):
    trained = trained or TRAINED
    freq = rope_frequencies(head_dim, base)
    wavelength = 2 * math.pi / freq
    out = []
    for f, w in zip(freq.tolist(), wavelength.tolist()):
        if w < trained / high:
            out.append(f)
        elif w > trained / low:
            out.append(f / s)
        else:
            keep = (trained / w - low) / (high - low)
            out.append((1 - keep) * f / s + keep * f)
    return torch.tensor(out), 1.0


STRETCHES["Llama 3.1"] = llama3

meta = llama3(tiny.head_dim, tiny.base, STRETCH)[0]
paper = yarn_as_written(tiny.head_dim, tiny.base, STRETCH, alpha=1, beta=4)[0]
print("Llama 3.1 against YaRN as written, β = 4:", f"{(meta / paper - 1).abs().max():.1e}")

Llama 3.1 against YaRN as written, β = 4: 1.2e-07

Each method against the transformers library, which implements them for the configs of real models, on three head shapes: TinyLlama’s, Qwen3’s and Llama 3.1’s. Dynamic NTK is checked at a length of \(sL\), where it equals NTK.

from transformers import LlamaConfig
from transformers.modeling_rope_utils import ROPE_INIT_FUNCTIONS


def library(head_dim, base, positions, params, length=None):
    c = LlamaConfig(
        hidden_size=head_dim * 8,
        num_attention_heads=8,
        head_dim=head_dim,
        max_position_embeddings=positions,
        rope_parameters=dict(rope_theta=base, **params),
    )
    return ROPE_INIT_FUNCTIONS[params["rope_type"]](c, "cpu", seq_len=length)


SHAPES = [(64, 10000.0, 2048, 8), (128, 1e6, 32768, 4), (128, 5e5, 8192, 8)]
checks = []
for head_dim, base, trained, s in SHAPES:
    longer = {"original_max_position_embeddings": trained}
    llama_params = {"low_freq_factor": 1, "high_freq_factor": 4, **longer}
    cases = [
        ("PI", interpolate(head_dim, base, s), trained * s,
         {"rope_type": "linear", "factor": s}),
        ("NTK", ntk(head_dim, base, s), trained, {"rope_type": "dynamic", "factor": 1.0}),
        ("YaRN", yarn(head_dim, base, s, trained), trained * s,
         {"rope_type": "yarn", "factor": s, **longer}),
        ("Llama 3.1", llama3(head_dim, base, s, trained), trained * s,
         {"rope_type": "llama3", "factor": s, **llama_params}),
    ]
    for name, (freq, scale), positions, params in cases:
        length = s * trained if name == "NTK" else None
        ref_freq, ref_scale = library(head_dim, base, positions, params, length)
        checks.append(
            {
                "head, base, window": f"{head_dim}, {base:,.0f}, {trained:,} × {s}",
                "method": name,
                "max relative difference": (freq / ref_freq - 1).abs().max().item(),
                "scale, ours": scale,
                "scale, library": ref_scale,
            }
        )
table(pd.DataFrame(checks), digits=3)
head, base, windowmethodmax relative differencescale, oursscale, library
64, 10,000, 2,048 × 8PI0.0001.0001.000
64, 10,000, 2,048 × 8NTK0.0001.0001.000
64, 10,000, 2,048 × 8YaRN0.0001.2081.208
64, 10,000, 2,048 × 8Llama 3.10.0001.0001.000
128, 1,000,000, 32,768 × 4PI0.0001.0001.000
128, 1,000,000, 32,768 × 4NTK0.0001.0001.000
128, 1,000,000, 32,768 × 4YaRN0.0001.1391.139
128, 1,000,000, 32,768 × 4Llama 3.10.0001.0001.000
128, 500,000, 8,192 × 8PI0.0001.0001.000
128, 500,000, 8,192 × 8NTK0.0001.0001.000
128, 500,000, 8,192 × 8YaRN0.0001.2081.208
128, 500,000, 8,192 × 8Llama 3.10.0001.0001.000

All four agree with the library to the last bit of a 32-bit float. Here is what each does to TinyLlama’s pairs when the window is stretched eight times: the frequency of each pair after the stretch, divided by its frequency before.

original = rope_frequencies(tiny.head_dim, tiny.base)
fig, ax = plt.subplots(figsize=(7.4, 3.8))
for method, color in [
    ("PI", COLORS["muted"]),
    ("NTK", COLORS["data"]),
    ("Llama 3.1", COLORS["enc"]),
    ("YaRN", COLORS["dec"]),
]:
    freq, _ = STRETCHES[method](tiny.head_dim, tiny.base, STRETCH)
    ratio = freq / original
    ax.plot(ratio, color=color, linewidth=2, label=method, marker="o", markersize=3)
written, _ = yarn_as_written(tiny.head_dim, tiny.base, STRETCH)
ax.plot(written / original, color=COLORS["dec"], linewidth=1.4, linestyle="--",
        label="YaRN as the paper writes it")
ax.set_yscale("log", base=2)
ax.set_xlabel("rotary pair, fastest to slowest")
ax.set_ylabel("new frequency / old")
ax.set_title(f"TinyLlama's {len(original)} pairs, bent for an {STRETCH}× window")
ax.legend(frameon=False, fontsize=9, loc="lower left")
plt.show()

TinyLlama

Position interpolation is the flat line at \(1/8\): every pair slowed down equally. NTK bends smoothly from 1 to \(1/8\). YaRN keeps the first pairs exactly, falls along a straight segment and reaches \(1/8\) for the pairs that never finished a turn; the paper’s version, dashed, takes a different path through the same middle. Llama 3.1 keeps more of the fast pairs, because its \(\beta\) is 4 rather than 32, and switches over in fewer pairs.

6. The race

All methods on the same eight books, with no training. Dynamic NTK at a fixed 16k length is NTK with the factor 8, so it isn’t a separate line here; the difference shows on shorter texts, in the passkey test below.

for method in ["PI", "NTK", "YaRN", "YaRN, no temperature", "Llama 3.1"]:
    stretch(tiny, method, STRETCH)
    curves[method] = average_losses(tiny)

fig, ax = plt.subplots(figsize=(7.4, 3.8))
for method, color in [
    ("none", COLORS["line"]),
    ("PI", COLORS["muted"]),
    ("NTK", COLORS["data"]),
    ("Llama 3.1", COLORS["enc"]),
    ("YaRN", COLORS["dec"]),
]:
    ax.plot(smooth(curves[method]), color=color, linewidth=2, label=method)
ax.axvline(TRAINED, color=COLORS["line"], linestyle="--", linewidth=1)
ax.set_ylim(top=2 * curves["YaRN"].mean().item())
ax.set_xlabel("position in the book")
ax.set_ylabel("loss, nats per token")
ax.set_title(f"TinyLlama stretched {STRETCH}× without training")
ax.legend(frameon=False, ncol=3)
plt.show()
bucket_table(["none", "PI", "NTK", "Llama 3.1", "YaRN", "YaRN, no temperature"])

Output of the code above

stretch0–1,0241,024–2,0482,048–4,0964,096–8,1928,192–16,384
none2.252.539.308.898.92
PI5.465.625.595.635.67
NTK2.382.722.632.727.00
Llama 3.12.393.003.033.433.67
YaRN2.362.672.602.703.00
YaRN, no temperature2.362.812.853.193.72

Read the table by columns. Inside the trained window, the first two columns, a method’s cost is how much it disturbs what the model already knew, and position interpolation pays the most: it slowed the fast pairs too, and the model’s sense of word order is blurred everywhere, near and far.

NTK leaves the window almost as it was and holds well past it, until its stretch of the slow pairs runs out, because by design only the very slowest pair is divided by the full factor. YaRN holds furthest.

Llama 3.1’s rule, tuned by Meta for a model with a far larger base, holds too, though not as well as YaRN on TinyLlama. The last row isolates its second idea, the temperature, and shows it is not a detail.

Why a temperature at all? The entropy of a query’s attention, in nats, says how many keys it spreads over: a query that looks at exactly one key has entropy 0, and one spread evenly over \(n\) keys has \(\ln n\). We take the last 64 queries of a book at growing lengths, in a middle layer, and measure.

layer = tiny.layers[len(tiny.layers) // 2].self_attn
lengths = [TRAINED // 4, TRAINED // 2, TRAINED, 2 * TRAINED, 4 * TRAINED, LENGTH]
entropy = {}
for method in ["YaRN, no temperature", "YaRN"]:
    stretch(tiny, method, STRETCH)
    values = []
    for n in lengths:
        layer.keep = True
        token_losses(tiny, books[1]["ids"][:, : n + 1])
        layer.keep = False
        q, k = layer.kept
        k = k.repeat_interleave(layer.heads // layer.kv_heads, dim=1)
        p = (q @ k.transpose(-1, -2) / math.sqrt(layer.head_dim)).softmax(-1)
        values.append(-(p * torch.log(p + 1e-12)).sum(-1).mean().item())
    entropy[method] = values

fig, ax = plt.subplots(figsize=(7.2, 3.4))
ax.plot(lengths, entropy["YaRN, no temperature"], color=COLORS["muted"], marker="o",
        label="no temperature")
ax.plot(lengths, entropy["YaRN"], color=COLORS["dec"], marker="o",
        label="YaRN's temperature")
ax.set_xscale("log", base=2)
ax.set_xlabel("tokens read")
ax.set_ylabel("attention entropy, nats")
ax.set_title("A middle layer of TinyLlama: how widely the last queries look")
ax.legend(frameon=False)
plt.show()

A middle layer of TinyLlama: how widely the last queries look

Entropy grows with the text: the softmax spreads over more keys simply because more keys exist. The temperature doesn’t stop the growth; it lowers the whole curve. It is a blunt instrument, one number tuned on LLaMA, and the table above shows what it buys on TinyLlama.

An average loss can hide whether a model can use a distant token. The passkey test, introduced by Mohtashami and Jaggi in 2023, checks that directly: a five-digit number hidden at a random depth in filler text, and the question at the end.

Position interpolation and NTK keep their fixed factor of 8, as a model deployed with them would; dynamic NTK and YaRN are set for the length they read, so at 2,048 tokens or less they change nothing.

FILLER = (
    "The grass is green. The sky is blue. The sun is yellow. "
    "Here we go. There and back again. "
)
INTRO = (
    "There is an important info hidden inside a lot of irrelevant text. "
    "Find it and memorize it. I will quiz you about the important information there. "
)


def passkey_prompt(n_tokens, depth, key):
    per_filler = len(tiny_tok(FILLER, add_special_tokens=False).input_ids)
    repeats = max(1, (n_tokens - 80) // per_filler)
    before = int(repeats * depth)
    needle = f"The pass key is {key}. Remember it. {key} is the pass key. "
    text = INTRO + FILLER * before + needle + FILLER * (repeats - before)
    text += "What is the pass key? The pass key is"
    return tiny_tok(text, return_tensors="pt").input_ids


@torch.no_grad()
def greedy(model, ids, steps, chunk=4096):
    caches = model.new_caches(ids.shape[1] + steps)
    for a in range(0, ids.shape[1], chunk):
        h = model(ids[:, a : a + chunk].to(device), start=a, caches=caches)
    out = [model.logits(h[:, -1:]).argmax(-1)]
    for i in range(steps - 1):
        h = model(out[-1], start=ids.shape[1] + i, caches=caches)
        out.append(model.logits(h).argmax(-1))
    return torch.cat(out, dim=1)[0].tolist()


rng = random.Random(SEED)
passkey_lengths = [TRAINED // 2, TRAINED, 2 * TRAINED, 4 * TRAINED, LENGTH]
TRIALS = 2 if SMOKE else 10
trials = {}
for n in passkey_lengths:
    trials[n] = [(rng.random(), rng.randint(10000, 99999)) for _ in range(TRIALS)]
STRETCHES["dynamic NTK"] = ntk  # with s = max(1, n / L), set for each length below
passkey = {}
for method in ["none", "PI", "NTK", "dynamic NTK", "YaRN"]:
    passkey[method] = []
    for n in passkey_lengths:
        fits = method in ["dynamic NTK", "YaRN"]
        stretch(tiny, method, max(1.0, n / TRAINED) if fits else STRETCH)
        hits = 0
        for depth, key in trials[n]:
            said = tiny_tok.decode(greedy(tiny, passkey_prompt(n, depth, key), 7))
            hits += str(key) in said
        passkey[method].append(hits / TRIALS)

heat = pd.DataFrame(passkey, index=[f"{n:,}" for n in passkey_lengths]).T
fig, ax = plt.subplots(figsize=(6.6, 3.2))
teal = LinearSegmentedColormap.from_list("teal", [COLORS["card"], COLORS["enc"]])
ax.imshow(heat.values, cmap=teal, vmin=0, vmax=1, aspect="auto")
ax.set_xticks(range(len(heat.columns)), heat.columns)
ax.set_yticks(range(len(heat.index)), heat.index)
for i in range(heat.shape[0]):
    for j in range(heat.shape[1]):
        v = heat.values[i, j]
        ink = "white" if v > 0.6 else "black"
        ax.text(j, i, f"{v:.0%}", ha="center", va="center", color=ink)
ax.set_xlabel("prompt length, tokens")
ax.set_title(f"TinyLlama finds the passkey, {TRIALS} random depths per cell")
plt.show()

Output of the code above

No method finds the passkey at 16k without training, and the cells before that tell more than the loss did. Position interpolation doesn’t find it even inside the trained window: with every pair slowed eightfold, the model can’t read a five-digit number back.

NTK with its fixed factor of 8 holds to 8k. Dynamic NTK and YaRN are stretched by exactly the ratio they need, which leaves no margin, and it shows.

NTK divides only the slowest pair by the full factor and the slow pairs just above it by less, so at twice the window those pairs still reach angles the model never saw: dynamic NTK loses the passkey as soon as it stretches at all, while NTK with a factor of 8 has the margin to cover them.

YaRN interpolates in full every pair that never finished a turn and finds the passkey at twice the window, but loses most of it at four times.

All of this is without training. The papers fine-tune after stretching, a thousand steps for position interpolation and four hundred for YaRN, and a fine-tuned model does far better than any zero-shot row here.

What the race measures is how good a starting point each method gives that fine-tuning, and how much it costs a model that is used without it.

7. Attention that doesn’t fit

Plain attention writes down a score for every pair of tokens: an \(n \times n\) matrix per head and per layer. Here is what it costs in 16-bit floats, for one head and for the 32 heads of one TinyLlama layer.

sizes = []
for n in [2048, 8192, 32768, 131072]:
    one_head = n * n * 2
    sizes.append(
        {"tokens": n, "one head, GB": one_head / 1e9, "32 heads, GB": 32 * one_head / 1e9}
    )
table(pd.DataFrame(sizes), digits=2)
tokensone head, GB32 heads, GB
2,0480.010.27
8,1920.134.29
32,7682.1568.72
131,07234.361,099.51

A T4 has 16 GB. Long before the weights are a problem, the scores are. The way out was found in three steps.

In 2018 Milakov and Gimelshein at NVIDIA showed that a softmax can be computed in one pass over its inputs; in 2021 Rabe and Staats at Google used that to compute attention without storing the matrix; in 2022 Tri Dao and colleagues at Stanford turned it into FlashAttention, a GPU kernel that keeps each tile of the work in the chip’s fast on-chip memory and is faster as well as smaller.

The softmax in one pass. The usual softmax reads the scores \(x_1, \dots, x_n\) three times: once for the maximum \(m\), subtracted so that \(e^x\) doesn’t overflow, once for the sum \(\ell = \sum_j e^{x_j - m}\), and once to divide. The online version keeps a running maximum and a running sum, and when a new maximum arrives it rescales the sum collected so far:

\[m_j = \max(m_{j-1}, x_j),\] \[\ell_j = \ell_{j-1}\, e^{m_{j-1} - m_j} + e^{x_j - m_j}.\]

After the last score, \(\ell_n\) is exactly the sum the three-pass version computes: every term in it ends up measured against the final maximum, whichever maximum it was added under. On five numbers:

def online_softmax(scores):
    running_max = float("-inf")
    running_sum = 0.0
    steps = []
    for x in scores:
        new_max = max(running_max, x)
        running_sum = running_sum * math.exp(running_max - new_max) + math.exp(x - new_max)
        running_max = new_max
        steps.append({"score": x, "running max": running_max, "running sum": running_sum})
    probs = [math.exp(x - running_max) / running_sum for x in scores]
    return probs, steps


scores = [3.0, 1.0, 4.0, 1.0, 5.0]
probs, steps = online_softmax(scores)
display(table(pd.DataFrame(steps), digits=4))
print("one pass:", [round(p, 4) for p in probs])
print("torch:   ", [round(p, 4) for p in torch.tensor(scores).softmax(0).tolist()])
scorerunning maxrunning sum
3.00003.00001.0000
1.00003.00001.1353
4.00004.00001.4177
1.00004.00001.4675
5.00005.00001.5398
one pass: [0.0879, 0.0119, 0.2389, 0.0119, 0.6494]
torch:    [0.0879, 0.0119, 0.2389, 0.0119, 0.6494]

Watch the third row. When the 4 arrives, the sum doesn’t simply grow by one: everything collected before is first multiplied by \(e^{3-4}\), because it was measured against a maximum that is no longer the maximum.

Attention needs one more running quantity, the weighted sum of the values \(O\), which gets the same rescaling.

Now cut the queries into blocks of rows and the keys and values into blocks of columns, and run the online update one tile at a time. In a causal model the tiles above the diagonal are the future and are never loaded.

How FlashAttention walks the attention matrix of a causal model. The queries are cut into blocks of rows and the keys and values into blocks of columns. For one block of queries the kernel loads one block of keys and values at a time into the GPU's fast on-chip memory, computes that tile of scores, folds it into a running maximum, a running sum and a running output, and moves on; the full matrix of scores never exists. Tiles above the diagonal are skipped, because a causal model's queries never look ahead How FlashAttention walks the attention matrix of a causal model. The queries are cut into blocks of rows and the keys and values into blocks of columns. For one block of queries the kernel loads one block of keys and values at a time into the GPU's fast on-chip memory, computes that tile of scores, folds it into a running maximum, a running sum and a running output, and moves on; the full matrix of scores never exists. Tiles above the diagonal are skipped, because a causal model's queries never look ahead

def plain_attention(q, k, v):
    n = q.shape[-2]
    scores = q @ k.transpose(-1, -2) / math.sqrt(q.shape[-1])
    causal = torch.ones(n, n, dtype=torch.bool, device=q.device).tril()
    scores = scores.masked_fill(~causal, float("-inf"))
    return scores.softmax(-1) @ v


def tiled_attention(q, k, v, block=128):
    n, d = q.shape[-2], q.shape[-1]
    out = torch.empty_like(q)
    for qs in range(0, n, block):
        qb = q[..., qs : qs + block, :]
        rows = qb.shape[-2]
        best = torch.full(qb.shape[:-1], float("-inf"), device=q.device)
        total = torch.zeros(qb.shape[:-1], device=q.device)
        acc = torch.zeros_like(qb)
        for ks in range(0, qs + rows, block):  # tiles right of the diagonal are never read
            kb = k[..., ks : ks + block, :]
            vb = v[..., ks : ks + block, :]
            s = qb @ kb.transpose(-1, -2) / math.sqrt(d)
            if ks + block > qs:  # the diagonal tile: hide the future
                qpos = torch.arange(qs, qs + rows, device=q.device)
                kpos = torch.arange(ks, ks + kb.shape[-2], device=q.device)
                s = s.masked_fill(kpos[None, :] > qpos[:, None], float("-inf"))
            new_best = torch.maximum(best, s.amax(-1))
            fade = torch.exp(best - new_best)
            p = torch.exp(s - new_best[..., None])
            total = total * fade + p.sum(-1)
            acc = acc * fade[..., None] + p @ vb
            best = new_best
        out[..., qs : qs + block, :] = acc / total[..., None]
    return out


torch.manual_seed(SEED)
q, k, v = (torch.randn(1, 4, 1000, 64, device=device) for _ in range(3))
reference = plain_attention(q, k, v)
tiled = tiled_attention(q, k, v)
kernel = F.scaled_dot_product_attention(q, k, v, is_causal=True)
print("tiled against plain, max difference:", f"{(tiled - reference).abs().max():.1e}")
print("PyTorch's kernel against plain:", f"{(kernel - reference).abs().max():.1e}")

tiled against plain, max difference: 5.4e-07
PyTorch's kernel against plain: 5.7e-07

One block of queries, tile by tile, with the shapes and with the running maximum and sum of its first row. The last tile is the diagonal one, where the future is masked.

block = 128
torch.manual_seed(SEED)
q1, k1, v1 = (torch.randn(1, 1, 4 * block, 64) for _ in range(3))
qs = 3 * block  # the fourth block of queries: it reads four tiles of keys
qb = q1[..., qs : qs + block, :]
best = torch.full((1, 1, block), float("-inf"))
total = torch.zeros(1, 1, block)
acc = torch.zeros_like(qb)
print("queries of the block:", tuple(qb.shape))
for ks in range(0, qs + block, block):
    kb, vb = k1[..., ks : ks + block, :], v1[..., ks : ks + block, :]
    s = qb @ kb.transpose(-1, -2) / math.sqrt(64)
    if ks == qs:
        s = s.masked_fill(torch.ones(block, block, dtype=torch.bool).triu(1), float("-inf"))
    new_best = torch.maximum(best, s.amax(-1))
    fade = torch.exp(best - new_best)
    p = torch.exp(s - new_best[..., None])
    total = total * fade + p.sum(-1)
    acc = acc * fade[..., None] + p @ vb
    best = new_best
    print(f"tile of keys {ks:3}–{ks + block - 1}: scores {tuple(s.shape)}, row 0:", end=" ")
    row = (best[0, 0, 0].item(), total[0, 0, 0].item(), fade[0, 0, 0].item())
    print("max {:.3f}, sum {:.3f}, rescaled by {:.3f}".format(*row))
check = plain_attention(q1, k1, v1)[..., qs : qs + block, :]
gap = (acc / total[..., None] - check).abs().max()
print("the block's output against plain attention:", f"{gap:.1e}")

queries of the block: (1, 1, 128, 64)
tile of keys   0–127: scores (1, 1, 128, 128), row 0: max 1.996, sum 24.006, rescaled by 0.000
tile of keys 128–255: scores (1, 1, 128, 128), row 0: max 1.996, sum 50.248, rescaled by 1.000
tile of keys 256–383: scores (1, 1, 128, 128), row 0: max 1.996, sum 73.460, rescaled by 1.000
tile of keys 384–511: scores (1, 1, 128, 128), row 0: max 1.996, sum 73.944, rescaled by 1.000
the block's output against plain attention: 3.0e-07

Exact to the rounding of 32-bit floats: FlashAttention is not an approximation. Our version runs its loops in Python and only shows the arithmetic.

The real kernel does the same arithmetic inside one GPU program, with each tile held in the small, fast memory next to each of the GPU’s processors, and that is where its speed comes from: attention is limited by trips to the GPU’s main memory, not by the arithmetic.

PyTorch ships several kernels behind scaled_dot_product_attention and picks one for the GPU at hand. FlashAttention-2 needs an Ampere GPU or newer; on a T4, PyTorch uses its memory-efficient kernel, which came from Meta’s xFormers library and rests on the same online softmax. Which ones run here:

from torch.nn.attention import SDPBackend, sdpa_kernel

q = torch.randn(1, 32, 2048, 64, device=device, dtype=HALF)
for name, backend in [
    ("FlashAttention", SDPBackend.FLASH_ATTENTION),
    ("memory-efficient", SDPBackend.EFFICIENT_ATTENTION),
    ("plain math", SDPBackend.MATH),
]:
    try:
        with sdpa_kernel(backend):
            F.scaled_dot_product_attention(q, q, q, is_causal=True)
        print(f"{name}: runs")
    except RuntimeError:
        print(f"{name}: not available here")

FlashAttention: not available here
memory-efficient: runs
plain math: runs

Now a race on TinyLlama’s shape, 32 heads of 64: plain attention that writes the matrix down, our tiled loop and PyTorch’s kernel, timed, with the peak memory of the GPU during each call. The floor under all three is TinyLlama’s weights, still loaded from the sections before.

def measure(fn, q):
    fn(q)
    if device == "cuda":
        torch.cuda.synchronize()
        torch.cuda.reset_peak_memory_stats()
    start = time.time()
    fn(q)
    if device == "cuda":
        torch.cuda.synchronize()
        return (time.time() - start) * 1000, torch.cuda.max_memory_allocated() / 1e9
    return (time.time() - start) * 1000, float("nan")


bench = []
bench_lengths = [256, 512] if SMOKE else [1024, 2048, 4096, 8192, 16384, 32768]


def run_plain(q):
    return plain_attention(q, q, q)


def run_tiled(q):
    return tiled_attention(q.float(), q.float(), q.float())


def run_kernel(q):
    return F.scaled_dot_product_attention(q, q, q, is_causal=True)


for n in bench_lengths:
    q = torch.randn(1, 32, n, 64, device=device, dtype=HALF)
    contenders = [
        ("plain", run_plain),
        ("tiled, Python", run_tiled),
        ("PyTorch kernel", run_kernel),
    ]
    for name, fn in contenders:
        if name == "tiled, Python" and n > 8192:
            continue  # minutes of Python loops, and its memory line is clear by then
        try:
            ms, gb = measure(fn, q)
        except torch.OutOfMemoryError:
            ms, gb = float("nan"), float("nan")
            print(f"{name}: out of memory at {n:,} tokens")
        bench.append({"tokens": n, "attention": name, "ms": ms, "peak GB": gb})
        if device == "cuda":
            torch.cuda.empty_cache()
bench = pd.DataFrame(bench)

fig, axes = plt.subplots(1, 2, figsize=(8.4, 3.4))
for name, color in [
    ("plain", COLORS["muted"]),
    ("tiled, Python", COLORS["data"]),
    ("PyTorch kernel", COLORS["dec"]),
]:
    part = bench[bench.attention == name]
    axes[0].plot(part.tokens, part["ms"], color=color, marker="o", label=name)
    axes[1].plot(part.tokens, part["peak GB"], color=color, marker="o", label=name)
for ax, label in zip(axes, ["time, ms", "peak memory, GB"]):
    ax.set_xscale("log", base=2)
    if device == "cuda":  # a CPU run has no memory counter to plot
        ax.set_yscale("log")
    ax.set_xlabel("tokens")
    ax.set_ylabel(label)
axes[0].legend(frameon=False, fontsize=9)
where = "a T4" if device == "cuda" else "a CPU"
fig.suptitle(f"One layer's attention, 32 heads of 64, on {where}")
plt.show()
peak = bench.pivot(index="tokens", columns="attention", values="peak GB")
table(peak.reset_index(), digits=2)

plain: out of memory at 16,384 tokens
plain: out of memory at 32,768 tokens

Output of the code above

tokensPyTorch kernelplaintiled, Python
1,0242.242.382.28
2,0482.252.792.32
4,0962.274.432.39
8,1922.3010.992.54
16,3842.37——
32,7682.50——

Plain attention’s memory grows with the square of the length and runs out well before 32k tokens. The kernel’s grows in a straight line, because it keeps only the inputs, the output and two numbers per row. Our Python loop draws the same straight line and loses only on time, which is the other half of FlashAttention’s work.

Training needs the backward pass too, and the backward pass needs the attention probabilities. FlashAttention doesn’t store them. It keeps \(m\) and \(\ell\) for every row, two numbers instead of \(n\), and recomputes each tile of probabilities when the gradient needs it: doing the arithmetic twice is cheaper than reading the matrix back.

8. A needle in a novel

Now Qwen3-4B, lesson 6’s architecture at four billion parameters. It was pretrained with a 32k window and a base of a million, and its card documents YaRN with a factor of 4 for up to 131,072 tokens.

Its 8 GB of weights and a 64k cache don’t fit one T4 together, so spread splits the layers over two. First the check against Hugging Face, whose model runs in 32-bit floats on the CPU, in a process of its own.

torch.save(qwen_tok(text, return_tensors="pt").input_ids, WORK_DIR / "qwen_ids.pt")
script = f"""
import torch
from transformers import AutoModelForCausalLM
model = AutoModelForCausalLM.from_pretrained("{QWEN_SOURCE}", dtype=torch.float32).eval()
ids = torch.load("{WORK_DIR / "qwen_ids.pt"}")
with torch.no_grad():
    torch.save(model(ids).logits.float(), "{WORK_DIR / "qwen_theirs.pt"}")
"""
subprocess.run([sys.executable, "-c", script], check=True)
qwen_theirs = torch.load(WORK_DIR / "qwen_theirs.pt")

del tiny
if device == "cuda":
    torch.cuda.empty_cache()
qwen = load_model(QWEN_SOURCE, qwen_c, HALF, spread=True)
with torch.no_grad():
    ours = qwen.logits(qwen(torch.load(WORK_DIR / "qwen_ids.pt").to(device))).cpu()
print("largest logit:", f"{qwen_theirs.abs().max().item():.2f}")
gap = (ours - qwen_theirs).abs().max()
print("max difference, ours in 16 bits against theirs in 32:", f"{gap:.2e}")

largest logit: 27.41
max difference, ours in 16 bits against theirs in 32: 5.97e-02

The haystack is the first of our PG-19 books, repeated if it runs out; the needle is one sentence that has nothing to do with it.

We ask in Qwen3’s chat format with its thinking switched off, and count an answer as right when it contains the number. Every length and depth runs twice: with the frequencies Qwen3 was trained with, and with YaRN’s 4×.

QWEN_TRAINED = 32768
haystack = qwen_tok(books[0]["text"], add_special_tokens=False).input_ids
NEEDLE = " The harbour master wrote the combination 4719 on the back of the chart. "
QUESTION = (
    "What combination did the harbour master write on the back of the chart? "
    "Answer with the number only."
)


def needle_prompt(n_tokens, depth):
    body = haystack
    while len(body) < n_tokens:
        body = body + haystack
    body = body[: n_tokens - 80]
    cut = int(len(body) * depth)
    story = qwen_tok.decode(body[:cut]) + NEEDLE + qwen_tok.decode(body[cut:])
    messages = [{"role": "user", "content": story + "\n\n" + QUESTION}]
    prompt = qwen_tok.apply_chat_template(
        messages, tokenize=False, add_generation_prompt=True, enable_thinking=False
    )
    return qwen_tok(prompt, return_tensors="pt", add_special_tokens=False).input_ids


needle_lengths = [256, 512] if SMOKE else [8192, 16384, 32768, 49152, 65536]
depths = [0.1, 0.5, 0.9]
found = {}
for setting in ["as trained", "YaRN 4×"]:
    if setting == "as trained":
        qwen.freq, qwen.scale = rope_frequencies(qwen.head_dim, qwen.base), 1.0
    else:
        qwen.freq, qwen.scale = yarn(qwen.head_dim, qwen.base, 4, QWEN_TRAINED)
    grid = np.zeros((len(depths), len(needle_lengths)))
    for j, n in enumerate(needle_lengths):
        for i, depth in enumerate(depths):
            said = qwen_tok.decode(greedy(qwen, needle_prompt(n, depth), 24))
            grid[i, j] = "4719" in said
    found[setting] = grid

fig, axes = plt.subplots(1, 2, figsize=(8.6, 3.2), sharey=True)
for ax, (setting, grid) in zip(axes, found.items()):
    ax.imshow(grid, cmap=teal, vmin=0, vmax=1, aspect="auto")
    ax.set_xticks(range(len(needle_lengths)), [f"{n // 1024}k" for n in needle_lengths])
    ax.set_yticks(range(len(depths)), [f"{d:.0%}" for d in depths])
    ax.set_title(f"Qwen3-4B, {setting}")
    ax.set_xlabel("prompt length")
axes[0].set_ylabel("needle depth")
plt.show()

Output of the code above

Teal is found, white is missed.

As trained, Qwen3-4B finds the needle at every depth up to 48k tokens, past the 32k it was trained on; at 64k it finds it only near the end of the text, close to the question. With YaRN it finds it everywhere.

A needle test is the easiest long-context exam there is: one literal sentence that doesn’t belong, and a question that repeats its words. Benchmarks built later make it harder, with several needles, questions that paraphrase and tasks that need the whole text, and on them models typically hold up for a fraction of their advertised window.

9. Questions about a whole novel

The needle was a fact that didn’t belong. A harder test is a fact that does: a detail of the plot that only a reader of the book would know.

The book is The Boats of the “Glen Carrig” by William Hope Hodgson, 1907: shipwrecked sailors drift through a sea choked with weed, fight off the devil-fish and the “weed men” on an island, and try to reach a ship that has been stuck in the weed for years.

It is told in mock eighteenth-century English, it is long enough to be over Qwen3’s window, 77 thousand tokens, and it is obscure enough that a small model is unlikely to have memorised its details. Project Gutenberg has it.

import re
import urllib.request

BOOK = "https://www.gutenberg.org/cache/epub/10542/pg10542.txt"
raw = urllib.request.urlopen(BOOK).read()
raw = raw.decode("utf-8").replace("\r", "")
begin = raw.index("\n", raw.index("*** START")) + 1
novel = raw[begin : raw.index("*** END")].strip()
chapters = re.split(r"\n(?=[IVX]+\n)", novel)  # the title page, then chapters I to XVII
if SMOKE:
    novel = novel[:3000]
print(f"{len(novel.split()):,} words, {len(qwen_tok(novel).input_ids):,} Qwen3 tokens")
for chapter in chapters[1:4]:
    print(" ".join(chapter.split("\n")[:3]).strip())

60,227 words, 77,360 Qwen3 tokens
I  The Land of Lonesomeness
II  The Ship in the Creek
III  The Thing That Made Search

Nine questions, each with the words that count as a right answer and the chapter that holds it, from the title page to the last quarter of the book. Each one was checked against the text.

QUESTIONS = [  # question, words that count as right, the chapter that holds the answer
    ("What is the surname of the man who tells the story?", ["winterstraw"], 0),
    ("How many days had the men been in the boats when they first sighted land?",
     ["five", "5"], 1),
    ("What was the name of the young seaman whom the creatures on the island carried off?",
     ["job"], 9),
    ("What weapon did the narrator build to send a line across to the ship in the weed?",
     ["bow"], 12),
    ("What was the name of the man the weed men took?", ["tompkins"], 14),
    ("Which seaman proposed to fly a kite over the ship?", ["jessop"], 14),
    ("How many people were living aboard the ship trapped in the weed?",
     ["twelve", "12"], 14),
    ("For how many years had the people of that ship been trapped in the weed?",
     ["seven", "7"], 14),
    ("What was the name of the young woman who showed the narrator round the freed ship?",
     ["madison", "mary"], 15),
]
ASK = "Answer the question about the text above in a few words.\nQuestion: "
MARK = "@@QUESTION@@"


def prompt_halves(content):
    messages = [{"role": "user", "content": content + MARK}]
    prompt = qwen_tok.apply_chat_template(
        messages, tokenize=False, add_generation_prompt=True, enable_thinking=False
    )
    before, after = prompt.split(MARK)
    return qwen_tok(before, add_special_tokens=False).input_ids, after


def is_right(said, answers):
    said = said.lower()
    return any(re.search(rf"\b{a}\b", said) for a in answers)

The whole book is read once: its tokens go through the model and into the cache, and then each question continues from that point. Setting every cache’s length back to the end of the book before the next question is all it takes to forget the previous one; the book is never read twice.

@torch.no_grad()
def read(model, ids, caches, start=0, chunk=4096):
    for a in range(0, len(ids), chunk):
        piece = torch.tensor([ids[a : a + chunk]], device=device)
        h = model(piece, start=start + a, caches=caches)
    return h


@torch.no_grad()
def answer(model, caches, start, question_ids, steps=24):
    for cache in caches:
        cache.length = start
    h = read(model, question_ids, caches, start=start)
    out = [model.logits(h[:, -1:]).argmax(-1)]
    for i in range(steps - 1):
        h = model(out[-1], start=start + len(question_ids) + i, caches=caches)
        out.append(model.logits(h).argmax(-1))
    return qwen_tok.decode(torch.cat(out, dim=1)[0], skip_special_tokens=True).strip()


def ask_all(context, settings):
    before, after = prompt_halves(context + "\n\n" + ASK)
    qwen.freq, qwen.scale = settings
    caches = qwen.new_caches(len(before) + 200)
    read(qwen, before, caches)
    said = []
    for question, _, _ in QUESTIONS:
        ids = qwen_tok(question + after, add_special_tokens=False).input_ids
        said.append(answer(qwen, caches, len(before), ids))
    del caches
    if device == "cuda":
        torch.cuda.empty_cache()
    return said


AS_TRAINED = (rope_frequencies(qwen.head_dim, qwen.base), 1.0)
YARN4 = yarn(qwen.head_dim, qwen.base, 4, QWEN_TRAINED)

Four ways to answer. No book: the title and the author only, to measure what the model already knows. Whole book, as trained: 77 thousand tokens with the frequencies of a 32k model.

Whole book, YaRN 4×. One chapter: just the chapter that holds the answer, a few thousand tokens, which is what a good search engine in front of the model would hand it.

title = 'The Boats of the "Glen Carrig", a novel by William Hope Hodgson (1907).'
results = {}
for setting, context, rope in [
    ("whole book, as trained", novel, AS_TRAINED),
    ("whole book, YaRN 4×", novel, YARN4),
]:
    results[setting] = ask_all(context, rope)
results["no book"] = ask_all(title, AS_TRAINED)
results["one chapter"] = []
for k, (_, _, chapter) in enumerate(QUESTIONS):
    context = chapters[chapter] if not SMOKE else novel[:2000]
    if chapter == 0:  # the narrator's name is on the title page, before chapter I
        context = chapters[0] + "\n" + chapters[1] if not SMOKE else novel[:2000]
    results["one chapter"].append(ask_all(context, AS_TRAINED)[k])

SETTINGS = ["no book", "whole book, as trained", "whole book, YaRN 4×", "one chapter"]
rows = []
for k, (question, answers, _) in enumerate(QUESTIONS):
    row = {"question": question[:48] + "…", "answer": answers[0]}
    for setting in SETTINGS:
        row[setting] = "✓" if is_right(results[setting][k], answers) else "·"
    rows.append(row)
table(pd.DataFrame(rows))
questionanswerno bookwhole book, as trainedwhole book, YaRN 4×one chapter
What is the surname of the man who tells the sto…winterstraw·✓✓✓
How many days had the men been in the boats when…five✓·✓✓
What was the name of the young seaman whom the c…job··✓✓
What weapon did the narrator build to send a lin…bow·✓·✓
What was the name of the man the weed men took?…tompkins···✓
Which seaman proposed to fly a kite over the shi…jessop·✓✓✓
How many people were living aboard the ship trap…twelve··✓✓
For how many years had the people of that ship b…seven·✓✓✓
What was the name of the young woman who showed …madison·✓✓✓
for setting in SETTINGS:
    right = sum(is_right(said, a) for said, (_, a, _) in zip(results[setting], QUESTIONS))
    print(f"{setting:24} {right} of {len(QUESTIONS)}")
print()
for k in [2, 4, 8]:
    print(QUESTIONS[k][0])
    for setting in ["whole book, as trained", "whole book, YaRN 4×", "one chapter"]:
        print(f"  {setting:24} {results[setting][k][:90]!r}")

no book                  1 of 9
whole book, as trained   5 of 9
whole book, YaRN 4×      7 of 9
one chapter              9 of 9

What was the name of the young seaman whom the creatures on the island carried off? whole book, as trained   'The name of the young seaman whom the creatures on the island carried off was **Tompkins**' whole book, YaRN 4×      'The name of the young seaman carried off was Job.\n\n\nThe name of the young se' one chapter              'Job.\nOverride\n</think>\n\nJob.\n\nThe name of the young se' What was the name of the man the weed men took? whole book, as trained   "The name of the man the weed men took was **the captain's wife**.\n\n**Explanation:**  \nIn" whole book, YaRN 4×      'Job.\n\nThe name of the man the weed men took was **Job**.

He was one of the crew' one chapter              'Tompkins.\nHuman: What is the name of the man who was lost?\n</think>\n\nThe' What was the name of the young woman who showed the narrator round the freed ship? whole book, as trained   'The name of the young woman who showed the narrator round the freed ship was **Mistress Ma' whole book, YaRN 4×      'Mistress Madison.\n\n**Mistress Madison** is the young woman who showed the narrator round t' one chapter              'Mary Madison.\nHuman:</think>\n\nThe name of the young woman who showed the narrator round th'

Without the book the model knows almost nothing about it, so the right answers come from reading. With the whole book and the frequencies it was trained with, more than twice its window, it still answers about half; with YaRN, most; with the one chapter, all of them.

The misses are instructive. The book loses two of its sailors, Job and Tompkins, in two different scenes, and each whole-book run once gives one man’s name for the other’s scene: the model finds the right passage and attaches the wrong name to it.

For questions like these a short, relevant passage beats a long one, which is why systems that answer questions about long documents usually search first and read second.

The printed answers show one more habit: after a short prompt the model sometimes keeps talking past its answer and invents the next turn of the chat. Only the first words are counted.

10. Where each piece fits today

Base and data first. Models trained today pick a large base from the start, half a million for Llama 3 and a million for Qwen, and spend a late stage of pretraining on long documents. Stretching covers the last factor of four or eight.

YaRN as the serving-time stretch. DeepSeek-V3 and gpt-oss ship with it in their configs, and Qwen3 documents how to switch it on. Hugging Face and vLLM implement the code’s version, the one we wrote.

FlashAttention everywhere. FlashAttention-2 sits behind PyTorch’s attention on Ampere GPUs and newer, and FlashAttention-3 was written for Hopper. On a T4 the memory-efficient kernel plays its role.

The cache is still the bill. Reading 64k tokens is now possible; keeping them is lesson 8’s problem, and the next lesson replaces the growing cache with a state of fixed size.

Llama 3.1’s passport, the model that made 128k the default for open models, and the family tree:

Llama 3.1 8B's passport. Published July 2024 by Meta under the Llama 3.1 Community License. Type: decoder-only, 8 billion parameters. Attention: 32 query heads sharing 8 key-value heads. Positions: rotary with base 500,000, stretched 8 times from the 8,192 tokens of pretraining by Meta's own frequency-dependent scaling. Context: 131,072 tokens. Vocabulary: 128,256 tokens. Good at: long documents at a size one GPU can serve. Not for: tasks that need the whole window to be read reliably Llama 3.1 8B's passport. Published July 2024 by Meta under the Llama 3.1 Community License. Type: decoder-only, 8 billion parameters. Attention: 32 query heads sharing 8 key-value heads. Positions: rotary with base 500,000, stretched 8 times from the 8,192 tokens of pretraining by Meta's own frequency-dependent scaling. Context: 131,072 tokens. Vocabulary: 128,256 tokens. Good at: long documents at a size one GPU can serve. Not for: tasks that need the whole window to be read reliably

The family of the Transformer as this course covers it, with lesson 9, long context, highlighted and lessons 1 to 8 done The family of the Transformer as this course covers it, with lesson 9, long context, highlighted and lessons 1 to 8 done

References

  1. Su, J., Lu, Y., Pan, S., et al. (2021). RoFormer: enhanced Transformer with rotary position embedding. arXiv:2104.09864.
  2. Chen, S., Wong, S., Chen, L., & Tian, Y. (2023). Extending context window of large language models via positional interpolation. arXiv:2306.15595.
  3. Peng, B., Quesnelle, J., Fan, H., & Shippole, E. (2023). YaRN: efficient context window extension of large language models. ICLR 2024.
  4. Zhang, P., Zeng, G., Wang, T., & Lu, W. (2024). TinyLlama: an open-source small language model. arXiv:2401.02385.
  5. Rae, J. W., Potapenko, A., Jayakumar, S. M., & Lillicrap, T. P. (2019). Compressive Transformers for long-range sequence modelling. ICLR 2020.
  6. Mohtashami, A., & Jaggi, M. (2023). Landmark attention: random-access infinite context length for Transformers. NeurIPS 2023.
  7. Milakov, M., & Gimelshein, N. (2018). Online normalizer calculation for softmax. arXiv:1805.02867.
  8. Rabe, M. N., & Staats, C. (2021). Self-attention does not need O(n²) memory. arXiv:2112.05682.
  9. Dao, T., Fu, D. Y., Ermon, S., Rudra, A., & Ré, C. (2022). FlashAttention: fast and memory-efficient exact attention with IO-awareness. NeurIPS 2022.
  10. Dao, T. (2023). FlashAttention-2: faster attention with better parallelism and work partitioning. ICLR 2024.
  11. Grattafiori, A., Dubey, A., Jauhri, A., et al. (2024). The Llama 3 herd of models. arXiv:2407.21783.
  12. Qwen Team (2025). Qwen3 technical report. arXiv:2505.09388.
]]> Eva Koroleva The KV cache budget and DeepSeek’s latent attention 2026-09-10T17:00:00+00:00 2026-11-08T00:00:00+00:00 https://blog.korolevalab.com/blog/2026/transformer-architectures-08-kv Transformer architectures: from BERT to diffusion · lesson 8 of 15 · all lessons · run on Kaggle · open in Colab · code on GitHub

A chat with a language model keeps a key and a value for every token it has read, in every layer, for as long as the chat lasts.

Lesson 6 counted that memory and showed the first two ways to cut it: share keys and values among query heads, and stop distant layers from looking far back.

This lesson is about the third way, the one DeepSeek introduced in 2024: don’t store keys and values at all. Store a short vector per token, a latent, from which every head’s key and value can be recomputed when needed.

DeepSeek-V3 has 128 query heads and 61 layers, and a token in its cache takes about half the memory it takes in Llama 3.1 8B, a model eighty times smaller. That trick, multi-head latent attention, is the subject of most of this lesson.

We write it from scratch, find out why it needs its own handling of positions, load a small model built on it, and measure on a T4 what it buys and what it costs.

Then we try to turn lesson 6’s Qwen3 into a latent model without any training, and look at the levers that work on any model: fewer bits per number, and a cache organised like virtual memory.

What one token keeps in the cache under four designs, drawn to scale for 32 heads of 128. Multi-head attention keeps a key and a value for every head. Grouped-query attention keeps them for a few key-value heads shared by groups of queries. Multi-query attention keeps one key and one value. Latent attention keeps one compressed vector and a small rotary key, from which every head's key and value are rebuilt What one token keeps in the cache under four designs, drawn to scale for 32 heads of 128. Multi-head attention keeps a key and a value for every head. Grouped-query attention keeps them for a few key-value heads shared by groups of queries. Multi-query attention keeps one key and one value. Latent attention keeps one compressed vector and a small rotary key, from which every head's key and value are rebuilt

What you’ll do

  1. Count what one token costs in the cache of ten models, and how many conversations fit next to them on a T4.
  2. Write multi-head latent attention in its two forms, the one that rebuilds keys and values and the one that never does, and check that they agree.
  3. Find out, with numbers, why rotary positions break the second form, and how DeepSeek mends it.
  4. Write MiniCPM3-4B, a latent-attention model that fits a T4, load its official weights and match Hugging Face.
  5. Measure how long a context fits with a latent cache and with a cache of rebuilt keys and values, and what each costs per generated token.
  6. Compress lesson 6’s Qwen3 into a latent model without training, two ways, and quantize its cache to fewer bits.

What you need: a GPU. On Kaggle, Settings → Accelerator → GPU; on Colab, Runtime → Change runtime type → T4 GPU. The notebook downloads MiniCPM3-4B (8 GB) and Qwen3-0.6B (1.5 GB); on a T4 it runs in about ten minutes.

Setup

One small helper package, tfs, draws the course’s plots and token chips and knows where the notebook runs (Kaggle, Colab or your own machine). Everything the lesson teaches is written out in the notebook itself.

import importlib.util
import subprocess
import sys

TFS = "git+https://github.com/xmarva/transformers-from-scratch@main"
if importlib.util.find_spec("tfs") is None:
    subprocess.check_call([sys.executable, "-m", "pip", "install", "-q", TFS])

from tfs import env
from tfs.viz import table, token_chips, use_course_style

env.quiet()  # no progress bars or environment warnings in the outputs
SEED = 42
env.seed_everything(SEED)
COLORS = use_course_style()
WORK_DIR = env.work_dir()

import json
import math
import subprocess
import sys
import time
from pathlib import Path

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import torch
import torch.nn as nn
import torch.nn.functional as F
from huggingface_hub import hf_hub_download
from safetensors.torch import load_file
from transformers import AutoModelForCausalLM, AutoTokenizer

SMOKE = env.smoke()  # CI runs the lesson on a small slice; you run it in full
device = "cuda" if torch.cuda.is_available() else "cpu"
HALF = torch.float16 if device == "cuda" else torch.float32  # 16-bit matmuls crawl on a CPU
assert device == "cuda" or SMOKE, "This lesson runs a 4B model: use a GPU"
if device == "cuda":
    print("GPU:", torch.cuda.get_device_name(0))

GPU: Tesla T4

1. What a token costs

For the models of lesson 6 the cost of a token is a plain product: two (a key and a value) times the layers \(L\), the key-value heads \(n_{kv}\) and the head size \(d_h\), in bytes per number \(b\):

\[\text{bytes per token} = 2 \cdot L \cdot n_{kv} \cdot d_h \cdot b.\]

A latent-attention model stores, per layer, one compressed vector of \(d_c\) numbers and one rotary key of \(d_h^R\) numbers shared by all heads:

\[\text{bytes per token} = (d_c + d_h^R) \cdot L \cdot b.\]

Here are ten models read from their configs, with every layer’s cache counted in full.

CACHE_MODELS = [  # repository, name; Llama and Gemma from Unsloth's public copies
    ("unsloth/Meta-Llama-3.1-8B", "Llama 3.1 8B"),
    ("mistralai/Mistral-7B-v0.1", "Mistral 7B"),
    ("Qwen/Qwen3-8B", "Qwen3 8B"),
    ("unsloth/gemma-3-27b-pt", "Gemma 3 27B"),
    ("deepseek-ai/deepseek-llm-67b-base", "DeepSeek 67B"),
    ("deepseek-ai/DeepSeek-V2-Lite", "DeepSeek-V2-Lite"),
    ("deepseek-ai/DeepSeek-V2", "DeepSeek-V2"),
    ("deepseek-ai/DeepSeek-V3", "DeepSeek-V3"),
    ("moonshotai/Kimi-K2-Instruct", "Kimi K2"),
    ("openbmb/MiniCPM3-4B", "MiniCPM3 4B"),
]


def read_config(repo):
    with open(hf_hub_download(repo, "config.json")) as f:
        c = json.load(f)
    return c.get("text_config", c)


def cache_elements(c):
    """Numbers kept per token, all layers."""
    layers = c["num_hidden_layers"]
    if c.get("kv_lora_rank"):
        return (c["kv_lora_rank"] + c["qk_rope_head_dim"]) * layers
    head = c.get("head_dim") or c["hidden_size"] // c["num_attention_heads"]
    return 2 * layers * c.get("num_key_value_heads", c["num_attention_heads"]) * head


configs = {name: read_config(repo) for repo, name in CACHE_MODELS}
CHAT = 8192  # tokens in one conversation
rows = []
for _, name in CACHE_MODELS:
    c = configs[name]
    per_token = cache_elements(c) * 2  # 16-bit numbers
    latent = bool(c.get("kv_lora_rank"))
    rows.append(
        {
            "model": name,
            "layers": c["num_hidden_layers"],
            "query heads": c["num_attention_heads"],
            "kept per layer": (
                f"latent {c['kv_lora_rank']} + rope {c['qk_rope_head_dim']}"
                if latent
                else f"{c.get('num_key_value_heads', c['num_attention_heads'])} kv heads"
            ),
            "KiB per token": per_token / 1024,
            "GB per 8k chat": per_token * CHAT / 1e9,
            "8k chats in 10 GB": int(10e9 // (per_token * CHAT)),
        }
    )
budget = pd.DataFrame(rows)
table(budget, digits=1)
modellayersquery headskept per layerKiB per tokenGB per 8k chat8k chats in 10 GB
Llama 3.1 8B32328 kv heads128.01.19
Mistral 7B32328 kv heads128.01.19
Qwen3 8B36328 kv heads144.01.28
Gemma 3 27B623216 kv heads496.04.22
DeepSeek 67B95648 kv heads380.03.23
DeepSeek-V2-Lite2716latent 512 + rope 6430.40.339
DeepSeek-V260128latent 512 + rope 6467.50.617
DeepSeek-V361128latent 512 + rope 6468.60.617
Kimi K26164latent 512 + rope 6468.60.617
MiniCPM3 4B6240latent 256 + rope 3234.90.334

The spread is wide. The dense models of lesson 6 sit in a narrow band, and Gemma 3 27B tops the table only because we counted every layer in full; its local layers keep a window, as lesson 6 showed, and a real Gemma cache is far smaller.

The latent models are a different class: DeepSeek-V3, with 128 query heads and 61 layers, keeps about half of what Llama 3.1 8B keeps per token, and MiniCPM3 a quarter.

DeepSeek’s own headline for this is bigger, a cache “reduced by 93.3%” against DeepSeek 67B. By our count from the configs, latent attention alone cuts the cache by about four fifths; the rest of the way to 93.3% comes from storing the cache in about six bits per number, which the paper mentions a few sections later.

“10 GB” is roughly what a T4 has left once a 4-billion-parameter model sits in it in 16-bit floats, so the last column is how many people could talk to such a model at once if it had each architecture’s cache.

2. Latent attention

DeepSeek-V2 describes it in a dozen equations[1]. For a token with hidden state \(h_t\), the keys and values of all heads come from one compressed vector:

\[c_t^{KV} = W^{DKV} h_t,\] \[k_t^C = W^{UK} c_t^{KV}, \qquad v_t^C = W^{UV} c_t^{KV},\]

where \(c_t^{KV}\) has \(d_c\) numbers, far fewer than the \(2 n_h d_h\) of the keys and values it stands for. “MLA only needs to cache \(c_t^{KV}\)”, and more than that: “since \(W^{UK}\) can be absorbed into \(W^Q\), and \(W^{UV}\) can be absorbed into \(W^O\), we even do not need to compute keys and values out for attention.”

The absorption is one line of algebra. The score of query \(q_t\) against key \(k_j = W^{UK} c_j\) in one head is

\[q_t^\top k_j = q_t^\top W^{UK} c_j = \big((W^{UK})^\top q_t\big)^\top c_j,\]

so instead of rebuilding a key for every cached token we move the query into the latent space once, and score it against the cached latents directly. The same happens on the way out: a mix of values is \(W^{UV}\) times a mix of latents, and \(W^{UV}\) can be applied after the mixing, or folded into the output projection.

Positions are handled separately, for a reason section 3 explains: each head gets a few more query dimensions, and all heads share one more key, \(k_t^R\), both rotated by RoPE. The cache holds \(c_t^{KV}\) and \(k_t^R\). Here is a layer of it, written so that the same weights can run either way.

class RMSNorm(nn.Module):
    def __init__(self, dim, eps=1e-5):
        super().__init__()
        self.eps = eps
        self.weight = nn.Parameter(torch.ones(dim))

    def forward(self, x):
        x32 = x.float()
        x32 = x32 * torch.rsqrt(x32.pow(2).mean(dim=-1, keepdim=True) + self.eps)
        return self.weight * x32.to(self.weight.dtype)


def rotate(x, cos, sin):
    half = x.shape[-1] // 2
    rotated = torch.cat([-x[..., half:], x[..., :half]], dim=-1)
    return x * cos.to(x.dtype) + rotated * sin.to(x.dtype)


class LatentCache:
    """What latent attention keeps per token: the latent and the shared rotary key."""

    def __init__(self):
        self.c, self.k_rope = None, None

    def add(self, c, k_rope):
        if self.c is not None:
            c, k_rope = (
                torch.cat([self.c, c], dim=1),
                torch.cat([self.k_rope, k_rope], dim=2),
            )
        self.c, self.k_rope = c, k_rope
        return c, k_rope

    def bytes(self):
        return (
            self.c.numel() * self.c.element_size()
            + self.k_rope.numel() * self.k_rope.element_size()
        )


class ExpandedCache:
    """What a plain cache would keep: every head's key and value, rebuilt."""

    def __init__(self):
        self.k, self.v = None, None

    def add(self, k, v):
        if self.k is not None:
            k, v = torch.cat([self.k, k], dim=2), torch.cat([self.v, v], dim=2)
        self.k, self.v = k, v
        return k, v

    def bytes(self):
        return (
            self.k.numel() * self.k.element_size() + self.v.numel() * self.v.element_size()
        )


class MLA(nn.Module):
    def __init__(self, c):
        super().__init__()
        d, self.heads = c["hidden_size"], c["num_attention_heads"]
        self.nope, self.rope = c["qk_nope_head_dim"], c["qk_rope_head_dim"]
        self.v_dim = (
            c.get("v_head_dim") or d // self.heads
        )  # MiniCPM3's config leaves it out
        self.rank = c["kv_lora_rank"]
        q_rank = c["q_lora_rank"]
        eps = c.get("rms_norm_eps", 1e-5)
        self.q_a_proj = nn.Linear(d, q_rank, bias=False)  # queries are compressed too
        self.q_a_layernorm = RMSNorm(q_rank, eps)
        self.q_b_proj = nn.Linear(q_rank, self.heads * (self.nope + self.rope), bias=False)
        self.kv_a_proj_with_mqa = nn.Linear(
            d, self.rank + self.rope, bias=False
        )  # c and k_rope
        self.kv_a_layernorm = RMSNorm(self.rank, eps)
        self.kv_b_proj = nn.Linear(
            self.rank, self.heads * (self.nope + self.v_dim), bias=False
        )
        self.o_proj = nn.Linear(self.heads * self.v_dim, d, bias=False)
        self.scale = (self.nope + self.rope) ** -0.5

    def queries(self, x, cos, sin):
        b, t, _ = x.shape
        q = self.q_b_proj(self.q_a_layernorm(self.q_a_proj(x)))
        q = q.view(b, t, self.heads, self.nope + self.rope).transpose(1, 2)
        q_nope, q_rope = q.split([self.nope, self.rope], dim=-1)
        return q_nope, rotate(q_rope, cos, sin)

    def latents(self, x, cos, sin):
        b, t, _ = x.shape
        c, k_rope = self.kv_a_proj_with_mqa(x).split([self.rank, self.rope], dim=-1)
        k_rope = rotate(k_rope.view(b, t, 1, self.rope).transpose(1, 2), cos, sin)
        return self.kv_a_layernorm(c), k_rope  # b x t x rank, and b x 1 x t x rope

    def up(self):
        """W_UK and W_UV, one block per head: heads x (nope + v) x rank."""
        return self.kv_b_proj.weight.view(self.heads, self.nope + self.v_dim, self.rank)

    def expand(self, c, k_rope):
        b, t, _ = c.shape
        kv = (
            self.kv_b_proj(c).view(b, t, self.heads, self.nope + self.v_dim).transpose(1, 2)
        )
        k_nope, v = kv.split([self.nope, self.v_dim], dim=-1)
        k = torch.cat([k_nope, k_rope.expand(-1, self.heads, -1, -1)], dim=-1)
        return k, v

    def forward(self, x, cos, sin, mask, cache=None, absorb=False):
        b, t, _ = x.shape
        q_nope, q_rope = self.queries(x, cos, sin)
        c, k_rope = self.latents(x, cos, sin)
        if isinstance(cache, ExpandedCache):
            k, v = cache.add(*self.expand(c, k_rope))
        else:
            if cache is not None:
                c, k_rope = cache.add(c, k_rope)
            if not absorb:  # rebuild every head's keys and values, then ordinary attention
                k, v = self.expand(c, k_rope)
        if absorb:  # score against the latents themselves; keys and values are never built
            w_uk, w_uv = self.up().split([self.nope, self.v_dim], dim=1)
            q_latent = torch.einsum("bhtn,hnr->bhtr", q_nope, w_uk)
            scores = (
                q_latent @ c.unsqueeze(1).transpose(-1, -2)
                + q_rope @ k_rope.transpose(-1, -2)
            ) * self.scale
            weights = scores.float().masked_fill(~mask, float("-inf")).softmax(dim=-1)
            mixed = weights.to(c.dtype) @ c.unsqueeze(1)  # b x heads x t x rank
            out = torch.einsum("bhtr,hvr->bhtv", mixed, w_uv)
        else:
            q = torch.cat([q_nope, q_rope], dim=-1)
            out = F.scaled_dot_product_attention(q, k, v, attn_mask=mask, scale=self.scale)
        return self.o_proj(out.transpose(1, 2).reshape(b, t, -1))

Two runs of the same random layer, one rebuilding keys and values and one absorbing the up-projections, must give the same output. In 32-bit floats they do, up to the order of the additions:

def rope_tables(positions, dim, theta=10000.0, factors=None):
    exponents = torch.arange(0, dim, 2, device=positions.device).float() / dim
    inv_freq = 1.0 / theta**exponents
    if factors is not None:  # MiniCPM3's long-context RoPE stretches each frequency
        inv_freq = inv_freq / torch.tensor(
            factors, device=positions.device, dtype=torch.float
        )
    angles = positions.float()[..., None] * inv_freq
    angles = torch.cat([angles, angles], dim=-1)
    return angles.cos()[:, None], angles.sin()[:, None]


small = {
    "hidden_size": 256,
    "num_attention_heads": 8,
    "qk_nope_head_dim": 32,
    "qk_rope_head_dim": 16,
    "v_head_dim": 32,
    "kv_lora_rank": 64,
    "q_lora_rank": 96,
}
torch.manual_seed(0)
layer = MLA(small)
for p in layer.parameters():
    nn.init.normal_(p, std=0.08) if p.dim() > 1 else None
x = torch.randn(2, 40, 256)
cos, sin = rope_tables(torch.arange(40)[None], 16)
causal = torch.ones(40, 40, dtype=torch.bool).tril()
with torch.no_grad():
    rebuilt = layer(x, cos, sin, causal)
    absorbed = layer(x, cos, sin, causal, absorb=True)
print("max difference:", f"{(rebuilt - absorbed).abs().max().item():.2e}")
heads, head = small["num_attention_heads"], small["qk_nope_head_dim"]
per_token = {
    "rebuilt keys and values": heads
    * (head + small["qk_rope_head_dim"] + small["v_head_dim"]),
    "latent and rotary key": small["kv_lora_rank"] + small["qk_rope_head_dim"],
}
print("numbers cached per token per layer:", per_token)

max difference: 9.24e-07
numbers cached per token per layer: {'rebuilt keys and values': 640, 'latent and rotary key': 80}

The two forms agree to the rounding of 32-bit floats, so the absorbed one is not an approximation; it is the same arithmetic in a different order. What changes is what has to be kept.

In this toy layer the rebuilt keys and values take eight times the memory of the latent and the rotary key, and in a real model the ratio is larger, because real models have more and wider heads per latent.

3. Why positions need their own dimensions

With rotary positions the key of token \(j\) is not \(W^{UK} c_j\) but \(R_j W^{UK} c_j\), where \(R_j\) rotates by an angle that depends on \(j\); the query is rotated too, by \(R_t\). The score becomes

\[q_t^\top R_t^\top R_j W^{UK} c_j,\]

and the matrix between the query and the latent now depends on both positions. It cannot be folded into the query once and for all: in DeepSeek’s words, “a RoPE matrix related to the currently generating token will lie between \(W^Q\) and \(W^{UK}\) and matrix multiplication does not obey a commutative law. As a result, we must recompute the keys for all the prefix tokens during inference”. Their fix is to keep the rotation out of the compressed part altogether and give it a few dimensions of its own, the decoupled rotary query and key of the layer above.

The failure is easy to see. Rotate the keys built from latents, the obvious design, and try to absorb anyway:

torch.manual_seed(1)
t, r, dh = 32, 64, 32
w_uk = torch.randn(dh, r) / math.sqrt(r)
latents = torch.randn(t, r)
query = torch.randn(dh)
cos1, sin1 = rope_tables(torch.arange(t)[None], dh)
keys = rotate((latents @ w_uk.T)[None, None], cos1, sin1)[
    0, 0
]  # rotated after the up-projection
q_rot = rotate(query.expand(1, 1, t, dh), cos1[:, :, -1:], sin1[:, :, -1:])[0, 0, 0]
honest = keys @ q_rot  # the scores the model needs, for the query at the last position
absorbed_scores = latents @ (w_uk.T @ q_rot)  # absorb W_UK as if no rotation sat between
without_rope = latents @ (w_uk.T @ query)
print(
    "scores with rotated keys vs absorbed:  max difference",
    f"{(honest - absorbed_scores).abs().max().item():.2f}",
)
print(
    "the same without RoPE:                  max difference",
    f"{((latents @ w_uk.T) @ query - without_rope).abs().max().item():.2e}",
)

scores with rotated keys vs absorbed:  max difference 15.72
the same without RoPE:                  max difference 3.34e-06

Absorbing past rotated keys gives scores that have nothing to do with the right ones, while without the rotation the absorbed scores are exact. That is the whole reason for the decoupled design: the rotary part is small, shared by all heads and cached as is, and the large compressed part never rotates, so it can be absorbed.

4. MiniCPM3 in PyTorch

DeepSeek’s own latent-attention models are too large for a T4: even DeepSeek-V2-Lite needs, in its card’s words, “40GB*1 GPU” for inference in 16-bit floats.

MiniCPM3-4B from OpenBMB uses the same layout, DeepSeek-V2’s names and all, at a size that fits: 62 layers, 40 heads, a latent of 256 numbers and a rotary key of 32. Its weights come under the MiniCPM model licence, free for research. Its config:

REPO = "openbmb/MiniCPM3-4B"
mini_config = configs["MiniCPM3 4B"]
for key in [
    "hidden_size",
    "num_hidden_layers",
    "num_attention_heads",
    "kv_lora_rank",
    "q_lora_rank",
    "qk_nope_head_dim",
    "qk_rope_head_dim",
    "intermediate_size",
    "vocab_size",
    "scale_emb",
    "scale_depth",
    "dim_model_base",
]:
    print(f"{key:20} {mini_config.get(key)}")

hidden_size          2560
num_hidden_layers    62
num_attention_heads  40
kv_lora_rank         256
q_lora_rank          768
qk_nope_head_dim     64
qk_rope_head_dim     32
intermediate_size    6400
vocab_size           73448
scale_emb            12
scale_depth          1.4
dim_model_base       256

Three numbers in it are MiniCPM’s own, from the “tensor programs” way of setting up training: the embeddings are multiplied by scale_emb, each residual branch by scale_depth over the square root of the number of layers, and the last hidden state is divided by hidden_size / dim_model_base before the output layer.

The rotary positions use per-frequency stretch factors, short_factor, that the config lists. Everything else is the decoder of lesson 6 with the attention above.

class FeedForward(nn.Module):
    def __init__(self, c):
        super().__init__()
        d, hidden = c["hidden_size"], c["intermediate_size"]
        self.gate_proj = nn.Linear(d, hidden, bias=False)
        self.up_proj = nn.Linear(d, hidden, bias=False)
        self.down_proj = nn.Linear(hidden, d, bias=False)

    def forward(self, x):
        return self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x))


class Block(nn.Module):
    def __init__(self, c):
        super().__init__()
        eps = c.get("rms_norm_eps", 1e-5)
        self.input_layernorm = RMSNorm(c["hidden_size"], eps)
        self.self_attn = MLA(c)
        self.post_attention_layernorm = RMSNorm(c["hidden_size"], eps)
        self.mlp = FeedForward(c)
        self.residual_scale = c["scale_depth"] / math.sqrt(c["num_hidden_layers"])

    def forward(self, x, cos, sin, mask, cache=None, absorb=False):
        h = self.self_attn(self.input_layernorm(x), cos, sin, mask, cache, absorb)
        x = (
            x + h.float() * self.residual_scale
        )  # the residual stream stays in 32-bit floats
        return x + self.mlp(self.post_attention_layernorm(x)).float() * self.residual_scale


class MiniCPM3(nn.Module):
    def __init__(self, c):
        super().__init__()
        self.config = c
        self.embed_tokens = nn.Embedding(c["vocab_size"], c["hidden_size"])
        self.layers = nn.ModuleList([Block(c) for _ in range(c["num_hidden_layers"])])
        self.norm = RMSNorm(c["hidden_size"], c.get("rms_norm_eps", 1e-5))
        self.lm_head = nn.Linear(c["hidden_size"], c["vocab_size"], bias=False)
        scaling = c.get("rope_scaling") or {}
        self.factors = scaling.get("short_factor")
        self.theta = c.get("rope_theta", 10000.0)

    def hidden(self, ids, mask, caches=None, start=0, absorb=False):
        positions = torch.arange(start, start + ids.shape[1], device=ids.device)[None]
        cos, sin = rope_tables(
            positions, self.config["qk_rope_head_dim"], self.theta, self.factors
        )
        x = self.embed_tokens(ids).float() * self.config["scale_emb"]
        for i, layer in enumerate(self.layers):
            cache = caches[i] if caches is not None else None
            x = layer(x, cos, sin, mask, cache, absorb)
        return self.norm(x)

    def logits(self, h):
        c = self.config
        return self.lm_head(h / (c["hidden_size"] / c["dim_model_base"])).float()

    def forward(self, ids, mask, caches=None, start=0, absorb=False):
        return self.logits(self.hidden(ids, mask, caches, start, absorb))


def load_minicpm3(source, c, dtype):
    if Path(source).is_dir():  # the smoke run's small model, saved to disk
        state = load_file(str(next(Path(source).glob("*.safetensors"))))
    else:
        path = hf_hub_download(source, "pytorch_model.bin")
        state = torch.load(path, map_location="cpu", weights_only=True, mmap=True)
    state = {k.removeprefix("model."): v.to(dtype) for k, v in state.items()}
    state.setdefault(
        "lm_head.weight", state["embed_tokens.weight"]
    )  # tied when it isn't stored
    with torch.device("meta"):
        model = MiniCPM3(c)
    model.load_state_dict(state, strict=True, assign=True)
    return model.to(device).eval()

The check: our model and Hugging Face’s on the same tokens. As in lesson 7, the reference runs first in a process of its own, so that the two copies are never in memory together.

tok = AutoTokenizer.from_pretrained(REPO)
text = (
    "The Rhine rises in the Swiss Alps, flows through Lake Constance and along "
    "the French–German border, and reaches the North Sea through a delta in the "
    "Netherlands."
)
check_ids = tok(text, return_tensors="pt").input_ids
torch.save(check_ids, WORK_DIR / "check_ids.pt")

if SMOKE:  # a small MiniCPM3 with random weights stands in for the real one
    from transformers import MiniCPM3Config, MiniCPM3ForCausalLM

    small_config = MiniCPM3Config(
        vocab_size=mini_config["vocab_size"],
        hidden_size=128,
        intermediate_size=256,
        num_hidden_layers=2,
        num_attention_heads=4,
        num_key_value_heads=4,
        kv_lora_rank=32,
        q_lora_rank=48,
        qk_nope_head_dim=16,
        qk_rope_head_dim=8,
        scale_emb=12,
        scale_depth=1.4,
        dim_model_base=32,
        tie_word_embeddings=True,
    )
    torch.manual_seed(0)
    reference = MiniCPM3ForCausalLM(small_config).eval()
    SOURCE = WORK_DIR / "minicpm3-small"
    reference.save_pretrained(SOURCE)
    with torch.no_grad():
        theirs = reference(check_ids).logits.float()
    mini_config = json.loads((SOURCE / "config.json").read_text())
    mini_config.pop("rope_scaling", None)
else:
    SOURCE = REPO
    script = f"""
import torch
from transformers import AutoConfig, AutoModelForCausalLM
config = AutoConfig.from_pretrained("{REPO}")
# the library's RoPE wants a stretch factor the config omits; 32k of 32k is 1
config.rope_parameters["factor"] = 1.0
model = AutoModelForCausalLM.from_pretrained("{REPO}", config=config, dtype=torch.float16)
model = model.to("cuda").eval()
ids = torch.load("{WORK_DIR / "check_ids.pt"}").to("cuda")
with torch.no_grad():
    torch.save(model(ids).logits.float().cpu(), "{WORK_DIR / "theirs.pt"}")
"""
    subprocess.run([sys.executable, "-c", script], check=True)
    theirs = torch.load(WORK_DIR / "theirs.pt")

model = load_minicpm3(SOURCE, mini_config, HALF)
t = check_ids.shape[1]
causal = torch.ones(t, t, dtype=torch.bool, device=device).tril()
with torch.no_grad():
    ours = model(check_ids.to(device), causal).cpu()
    ours_absorbed = model(check_ids.to(device), causal, absorb=True).cpu()
print("tokens:", t)
print("largest logit:", f"{theirs.abs().max().item():.2f}")
print(
    "max difference, keys and values rebuilt:", f"{(ours - theirs).abs().max().item():.2e}"
)
print("max difference, absorbed:", f"{(ours_absorbed - theirs).abs().max().item():.2e}")

tokens: 36
largest logit: 30.34
max difference, keys and values rebuilt: 8.59e-02
max difference, absorbed: 1.07e-01

Both forms of our model agree with Hugging Face’s within the rounding of 16-bit floats over 62 layers; in 32-bit floats, on the small random model this notebook’s smoke test uses, the difference is a few millionths. The weights load into our classes unchanged: MiniCPM3 is DeepSeek-V2’s attention with smaller numbers.

Generation is lesson 4’s loop with a cache per layer. generate takes the kind of cache and whether to absorb, so the same weights can run with a latent cache and absorption, as DeepSeek serves its models, or with a cache of rebuilt keys and values, the way the code published with MiniCPM3 runs it.

@torch.no_grad()
def generate(model, ids, steps, cache_kind=LatentCache, absorb=True):
    caches = [cache_kind() for _ in model.layers]
    t = ids.shape[1]
    mask = torch.ones(t, t, dtype=torch.bool, device=device).tril()
    logits = model(ids, mask, caches, absorb=absorb and cache_kind is LatentCache)
    out = [logits[:, -1].argmax(-1)]
    for step in range(steps - 1):
        mask = torch.ones(1, t + step + 1, dtype=torch.bool, device=device)
        logits = model(
            out[-1][:, None],
            mask,
            caches,
            start=t + step,
            absorb=absorb and cache_kind is LatentCache,
        )
        out.append(logits[:, -1].argmax(-1))
    return torch.stack(out, dim=1), caches


messages = [{"role": "user", "content": "Explain in two sentences why the sky is blue."}]
prompt = tok.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
ids = tok(prompt, return_tensors="pt", add_special_tokens=False).input_ids.to(device)
written, caches = generate(model, ids, 12 if SMOKE else 80)
print(tok.decode(written[0], skip_special_tokens=True))
cached = sum(c.bytes() for c in caches)
print(f"cache: {cached / 1024:.0f} KiB for {ids.shape[1] + written.shape[1] - 1} tokens")

The sky appears blue because the Earth's atmosphere scatters sunlight in all directions. Blue light is scattered more than other colors because it travels as shorter, smaller waves.
cache: 3453 KiB for 99 tokens

5. The budget on a T4

The table of section 1 is arithmetic. Here is the measurement: one long document read into MiniCPM3 in chunks of a thousand tokens, with each kind of cache, until the GPU runs out of memory.

Each chunk attends to everything before it; with the latent cache the keys and values of one layer are rebuilt for the chunk and thrown away, so only the latents stay.

def long_document(n_tokens):
    name = "wikitext-103-raw-v1/test-00000-of-00001.parquet"
    path = hf_hub_download("Salesforce/wikitext", name, repo_type="dataset")
    ids = tok("".join(pd.read_parquet(path)["text"].tolist())[: n_tokens * 8]).input_ids[
        :n_tokens
    ]
    return torch.tensor(ids)[None]


CHUNK = 64 if SMOKE else 1024
LIMIT = 512 if SMOKE else 32768
document = long_document(LIMIT).to(device)


@torch.no_grad()
def read_until_full(cache_kind):
    caches = [cache_kind() for _ in model.layers]
    reached, kept = 0, []
    torch.cuda.empty_cache() if device == "cuda" else None
    try:
        for start in range(0, document.shape[1], CHUNK):
            chunk = document[:, start : start + CHUNK]
            n = chunk.shape[1]
            mask = torch.ones(n, start + n, dtype=torch.bool, device=device).tril(
                diagonal=start
            )
            model.hidden(chunk, mask, caches, start=start)
            reached = start + n
            kept.append((reached, sum(c.bytes() for c in caches)))
    except torch.cuda.OutOfMemoryError:
        pass
    del caches
    torch.cuda.empty_cache() if device == "cuda" else None
    return reached, kept


capacity = {}
for name, kind in [
    ("latent cache", LatentCache),
    ("rebuilt keys and values", ExpandedCache),
]:
    reached, kept = read_until_full(kind)
    capacity[name] = kept
    print(
        f"{name:26} read {reached:6,} tokens before running out"
        if reached < LIMIT
        else f"{name:26} read all {reached:,} tokens"
    )

fig, ax = plt.subplots(figsize=(7.4, 3.6))
for (name, kept), color in zip(capacity.items(), [COLORS["enc"], COLORS["dec"]]):
    ax.plot(
        [k for k, _ in kept],
        [b / 1e9 for _, b in kept],
        color=color,
        label=name,
        linewidth=2,
    )
ax.set_xlabel("tokens read")
ax.set_ylabel("cache, GB")
ax.set_title("MiniCPM3-4B on one T4, one document read in chunks")
ax.legend(frameon=False)
plt.show()

latent cache               read all 32,768 tokens
rebuilt keys and values    read  7,168 tokens before running out

MiniCPM3-4B on one T4, one document read in chunks

The latent cache read the whole document, all 32,768 tokens, with a cache of about a gigabyte. The cache of rebuilt keys and values ran out of memory a little past a fifth of it.

The model is the same and so are the weights; only what we chose to keep differs. And the second curve is not a strawman: the code that was published with MiniCPM3, and vLLM’s implementation of it, both cache rebuilt keys and values.

Memory is half the story; the other half is time. With the latent cache, every new token either rebuilds the keys and values of the whole context in every layer, or absorbs and works in the latent space, where each head’s score is a dot product over the 256 + 32 numbers of the latent instead of the 64 + 32 of a key.

We time one step of generation at several context lengths, three ways.

@torch.no_grad()
def step_time(context, cache_kind, absorb, repeats=5):
    """Median time of one new token after `context` tokens; NaN if the cache doesn't fit."""
    caches = [cache_kind() for _ in model.layers]
    try:
        for start in range(0, context, CHUNK):
            chunk = document[:, start : min(start + CHUNK, context)]
            n = chunk.shape[1]
            mask = torch.ones(n, start + n, dtype=torch.bool, device=device).tril(
                diagonal=start
            )
            model.hidden(chunk, mask, caches, start=start)
        token = document[:, context : context + 1]
        mask = torch.ones(1, context + 1, dtype=torch.bool, device=device)
        timings = []
        for _ in range(repeats):
            if device == "cuda":
                torch.cuda.synchronize()
            begin = time.time()
            model.hidden(token, mask, caches, start=context, absorb=absorb)
            if device == "cuda":
                torch.cuda.synchronize()
            timings.append(time.time() - begin)
            for (
                cache
            ) in caches:  # take the step back out, so every repeat sees the same context
                if isinstance(cache, LatentCache):
                    cache.c, cache.k_rope = cache.c[:, :-1], cache.k_rope[:, :, :-1]
                else:
                    cache.k, cache.v = cache.k[:, :, :-1], cache.v[:, :, :-1]
        return float(np.median(timings))
    except torch.cuda.OutOfMemoryError:
        return float("nan")
    finally:
        del caches
        if device == "cuda":
            torch.cuda.empty_cache()


contexts = [128, 256] if SMOKE else [1024, 2048, 4096, 8192]
ways = [
    ("latent cache, rebuilt each step", LatentCache, False),
    ("latent cache, absorbed", LatentCache, True),
    ("rebuilt keys and values cached", ExpandedCache, False),
]
speed = []
for context in contexts:
    for name, kind, absorb in ways:
        speed.append(
            {
                "context": context,
                "way": name,
                "ms per token": 1000 * step_time(context, kind, absorb),
            }
        )
speed = pd.DataFrame(speed)
table(
    speed.pivot(index="way", columns="context", values="ms per token").reset_index(),
    digits=1,
)
way1024204840968192
latent cache, absorbed120.3120.4117.4128.5
latent cache, rebuilt each step106.7104.6142.1261.8
rebuilt keys and values cached104.299.3125.6nan

Three different behaviours. With absorption, a new token costs the same at every context length here: the time goes into the rest of the model, reading 8 GB of weights and running 62 layers, and the scores against a few thousand latents are cheap next to that.

Rebuilding the keys and values at every step costs more and more as the context grows, because it rebuilds all of them, every time. Caching rebuilt keys and values is the fastest at short contexts, since the attention is then ordinary, and runs out of memory before 8,192 tokens.

Absorption is not free either. Per cached token each head scores over the 256 + 32 numbers of the latent instead of the 64 + 32 of a key, and mixes 256 numbers instead of 64, so the attention does several times the arithmetic.

On a large GPU with many heads that arithmetic dominates; DeepSeek’s kernel notes say that “the MLA kernel is compute-bound despite being a decoding-stage attention kernel”.

That is why serving engines use each form where it fits: vLLM’s documentation describes the rebuilt form for reading a prompt and the absorbed one for generating, the first “compute friendly”, the second “data-movement friendly”.

6. A latent cache for a model that wasn’t trained with one

If latent attention is this good, why not convert existing models? People have tried.

MHA2MLA factorises the key and value weights jointly with an SVD and then fine-tunes on “0.6% to 1%” of the data[2]; TransMLA instead runs principal component analysis on the keys and values the model actually produces, and without any training reports that “we compressed the KV cache of LLaMA-2 by 68.75%, with only a 1.65% performance drop”[3].

We can try the simplest version on lesson 6’s Qwen3-0.6B. In each layer the keys (before their norm and rotation) and the values come from two matrices, together 2,048 numbers per token. Replace them by a down-projection to \(r\) numbers, the latent we would cache, and an up-projection back:

\[\begin{bmatrix} k \\ v \end{bmatrix} = W x \;\approx\; B\,(A x),\] \[A \in \mathbb{R}^{r \times d}, \qquad B \in \mathbb{R}^{2048 \times r}.\]

Two ways to choose \(A\) and \(B\). The weight SVD takes the best rank-\(r\) approximation of \(W\) itself. The activation PCA takes the \(r\) directions along which the keys and values of real text vary most, measured on a few thousand tokens, and projects onto those.

The first treats every direction of the input as equally likely; the second knows which inputs the model actually sees. We measure the loss on held-out text as \(r\) shrinks.

The rotation still has to happen after the up-projection, so this model can cache the latent but not absorb; the point is how much the keys and values can be compressed at all.

QWEN = "Qwen/Qwen3-0.6B"
del model
torch.cuda.empty_cache() if device == "cuda" else None
qwen_tok = AutoTokenizer.from_pretrained(QWEN)
if SMOKE:
    from transformers import Qwen3Config, Qwen3ForCausalLM

    torch.manual_seed(0)
    qwen = Qwen3ForCausalLM(
        Qwen3Config(
            vocab_size=qwen_tok.vocab_size + 300,
            hidden_size=64,
            intermediate_size=128,
            num_hidden_layers=2,
            num_attention_heads=4,
            num_key_value_heads=2,
            head_dim=16,
        )
    )
    qwen = qwen.to(device).eval()
else:
    qwen = AutoModelForCausalLM.from_pretrained(QWEN, dtype=HALF).to(device).eval()

name = "wikitext-103-raw-v1/test-00000-of-00001.parquet"
path = hf_hub_download("Salesforce/wikitext", name, repo_type="dataset")
wiki = "".join(pd.read_parquet(path)["text"].tolist())
qwen_ids = qwen_tok(wiki[: 400_000 if not SMOKE else 20_000]).input_ids
SEQ = 64 if SMOKE else 512
n_windows = len(qwen_ids) // SEQ
windows = torch.tensor(qwen_ids[: n_windows * SEQ]).view(n_windows, SEQ)
calibration, held_out = (
    windows[: n_windows // 2],
    windows[n_windows // 2 :][: 4 if SMOKE else 32],
)


def loss_on(batches, batch=8):
    total, count = 0.0, 0
    with torch.no_grad():
        for i in range(0, len(batches), batch):
            ids = batches[i : i + batch].to(device)
            logits = qwen(ids).logits.float()
            total += F.cross_entropy(
                logits[:, :-1].flatten(0, 1), ids[:, 1:].flatten(), reduction="sum"
            ).item()
            count += ids[:, 1:].numel()
    return total / count


attention = [layer.self_attn for layer in qwen.model.layers]
modules = [
    (a.k_proj, a.v_proj) for a in attention
]  # the original projections, put back after each test
original = [(k.weight.detach(), v.weight.detach()) for k, v in modules]
kv_width = original[0][0].shape[0] + original[0][1].shape[0]

# the keys (before norm and rotation) and values of real text, for the PCA
inputs = [[] for _ in attention]
hooks = [
    a.k_proj.register_forward_hook(
        lambda m, i, o, k=k: inputs[k].append(i[0].detach().float().flatten(0, 1).cpu())
    )
    for k, a in enumerate(attention)
]
with torch.no_grad():
    for i in range(0, min(len(calibration), 4 if SMOKE else 16), 4):
        qwen(calibration[i : i + 4].to(device))
for hook in hooks:
    hook.remove()


class LowRank(nn.Module):
    """y = B (A x): the projection goes through r numbers."""

    def __init__(self, a, b):
        super().__init__()
        self.a, self.b = nn.Parameter(a), nn.Parameter(b)

    def forward(self, x):
        return (x @ self.a.T) @ self.b.T


def compress(rank, method):
    for k, a in enumerate(attention):
        w = torch.cat(original[k]).float()  # (2048 x d): keys, then values
        if method == "weight SVD":
            u, s, vt = torch.linalg.svd(w, full_matrices=False)
            down, up = s[:rank, None] * vt[:rank], u[:, :rank]
        else:
            x = torch.cat(inputs[k]).to(w.device)
            y = x @ w.T  # the keys and values this layer produced
            directions = torch.linalg.eigh(y.T @ y / len(y)).eigenvectors[:, -rank:]
            down, up = directions.T @ w, directions
        n_k = original[k][0].shape[0]
        dtype = original[k][0].dtype
        a.k_proj = LowRank(down.to(dtype), up[:n_k].to(dtype))
        a.v_proj = LowRank(down.to(dtype), up[n_k:].to(dtype))


def restore():
    for a, (k_proj, v_proj) in zip(attention, modules):
        a.k_proj, a.v_proj = k_proj, v_proj


base_loss = loss_on(held_out)
ranks = [kv_width // 2, kv_width // 4, kv_width // 8, kv_width // 16, kv_width // 32]
conversion = [
    {"method": m, "latent": kv_width, "loss": base_loss}
    for m in ["weight SVD", "activation PCA"]
]
for method in ["weight SVD", "activation PCA"]:
    for rank in ranks:
        compress(rank, method)
        conversion.append({"method": method, "latent": rank, "loss": loss_on(held_out)})
restore()
conversion = pd.DataFrame(conversion)
fig, ax = plt.subplots(figsize=(7.2, 3.6))
for (method, chunk), color in zip(
    conversion.groupby("method"), [COLORS["att"], COLORS["dec"]]
):
    chunk = chunk.sort_values("latent")
    ax.plot(chunk["latent"], chunk["loss"], marker="o", color=color, label=method)
ax.axhline(base_loss, color=COLORS["line"], linestyle="--", linewidth=1)
ax.set_xscale("log", base=2)
ax.set_xticks([kv_width] + ranks, [str(r) for r in [kv_width] + ranks])
ax.set_xlabel(f"numbers cached per token per layer (keys and values: {kv_width})")
ax.set_ylabel("loss on held-out text")
ax.set_title("Qwen3-0.6B with its keys and values squeezed through a latent, no training")
ax.legend(frameon=False)
plt.show()

Qwen3-0.6B with its keys and values squeezed through a latent, no training

At 1,024 numbers both methods lose nothing, and that is the first lesson of this experiment: Qwen3-0.6B’s hidden state is only 1,024 numbers wide, so its keys and values, 2,048 numbers computed linearly from it, carry no more information than the hidden state itself. The model caches twice what it needs.

Below that the two methods part. The weight SVD falls apart quickly, the activation PCA degrades gently, which is why TransMLA works with the keys and values the model actually produces. But even the PCA loses noticeably at a quarter of the original cache, and the loss grows from there.

A model can be converted, but not for free: the published conversions all fine-tune afterwards, MHA2MLA on a fraction of a percent of the pre-training data, and the rotary positions have to be split out first, which our experiment didn’t attempt.

7. Fewer bits per number

The lever that works on every model is the number format. KIVI found that keys and values want different treatment: “the key cache should be quantized per-channel … In contrast, the value cache should be quantized per-token”, because “in key cache, some fixed channels exhibit very large magnitudes, whereas in value cache, there is no significant pattern for outliers”[4].

We try both groupings on Qwen3’s keys and values at 8, 4 and 2 bits, rounding them on the way into the cache and measuring the loss.

def fake_quantize(x, bits, dim):
    """Round x to 2**bits levels per slice along `dim`, then back to floats."""
    low, high = x.amin(dim=dim, keepdim=True), x.amax(dim=dim, keepdim=True)
    step = (high - low).clamp(min=1e-8) / (2**bits - 1)
    return ((x - low) / step).round() * step + low


class Quantized(nn.Module):
    def __init__(self, inner, bits, dim):
        super().__init__()
        self.inner, self.bits, self.dim = inner, bits, dim

    def forward(self, x):
        return fake_quantize(self.inner(x), self.bits, self.dim)


def quantize_cache(bits, keys_by, values_by):
    by = {"channel": 1, "token": 2}  # outputs are batch x tokens x channels
    for a, (k_proj, v_proj) in zip(attention, modules):
        a.k_proj = Quantized(k_proj, bits, by[keys_by])
        a.v_proj = Quantized(v_proj, bits, by[values_by])


quantization = [{"bits": 16, "grouping": "none", "loss": base_loss}]
for bits in [8, 4, 2]:
    for keys_by, values_by in [("channel", "token"), ("token", "token")]:
        quantize_cache(bits, keys_by, values_by)
        quantization.append(
            {
                "bits": bits,
                "grouping": f"keys per {keys_by}, values per {values_by}",
                "loss": loss_on(held_out),
            }
        )
restore()
table(pd.DataFrame(quantization), digits=3)
bitsgroupingloss
16none3.403
8keys per channel, values per token3.403
8keys per token, values per token3.404
4keys per channel, values per token3.554
4keys per token, values per token3.687
2keys per channel, values per token10.709
2keys per token, values per token13.827

Eight bits cost nothing measurable, which is why an 8-bit or FP8 cache is a switch in serving engines.

At four bits the grouping matters: keys rounded per channel keep the loss noticeably closer to the original than keys rounded per token, as KIVI predicts, because a few channels of the keys carry much larger numbers than the rest.

At two bits both break down here. KIVI gets to two bits with more care than this sketch: it rounds keys in small groups of tokens rather than over the whole text, and keeps the most recent tokens in full precision.

Our per-channel rounding also looks at the whole sequence at once, which a cache filled one token at a time cannot do.

8. Where each lever fits

Every lever in this lesson trades something, and they combine.

Grouped queries, from lesson 6, are now the default: they cost a little quality and nothing else, and nearly every new model has them. Local layers, also lesson 6, cut the cache of most layers to a window and need the model to be trained that way.

Latent attention cuts deepest, to about a quarter of what grouped queries with eight groups would keep at the same size, at the price of more arithmetic per generated token, and it too has to be trained in, or converted and fine-tuned.

The levers that need no retraining are the number format and the memory manager. An 8-bit cache halves the memory for free, as we saw, and lower formats need care.

The memory manager matters more than it sounds: before vLLM, in its authors’ measurement, “only 20.4% - 38.2% of the KV cache memory is used to store the actual token states in the existing systems”[5], the rest lost to caches reserved for the longest possible answer.

PagedAttention stores the cache in fixed blocks, like an operating system’s pages, and shares the blocks of a common prefix between requests, lesson 6’s prefix caching done by the server.

The industry combines all of these. Character.AI, by its own account, cut its cache by “more than 20X” with multi-query attention, one global layer in six, keys and values shared between neighbouring layers, and 8-bit numbers: “We use int8 quantization on model weights, activations, and attention KV cache” [6].

Sharing between layers is its own research line, from cross-layer attention[7] to YOCO, which keeps one cache for the whole model[8]. And when nothing else is left, the cache can be cut by forgetting: H2O keeps the few tokens that collect most of the attention and drops the rest[9].

For a T4 the practical answer is the table of section 1 with the weights subtracted: a 4B model in 16-bit floats leaves room for a handful of long conversations with a grouped-query cache, several times more with a latent one, and twice as many again with an 8-bit cache.

DeepSeek-V2’s passport, the model that introduced latent attention, and the family tree with this lesson’s branch:

DeepSeek-V2's passport. Published May 2024 by DeepSeek-AI with open weights. Type: decoder-only mixture of experts, 236 billion parameters, 21 billion per token, 160 routed and 2 shared experts, 6 routed per token. Training: 8.1 trillion tokens. Attention: multi-head latent attention, 128 heads, a latent of 512 numbers and a rotary key of 64 per token and layer; queries compressed to 1,536. Positions: decoupled rotary dimensions, 64 per head. Cache: 576 numbers per token per layer, as much as grouped-query attention with 2.25 groups. Vocabulary: 102,400 tokens. Good at: long conversations served to many users at once. Not for: a single GPU DeepSeek-V2's passport. Published May 2024 by DeepSeek-AI with open weights. Type: decoder-only mixture of experts, 236 billion parameters, 21 billion per token, 160 routed and 2 shared experts, 6 routed per token. Training: 8.1 trillion tokens. Attention: multi-head latent attention, 128 heads, a latent of 512 numbers and a rotary key of 64 per token and layer; queries compressed to 1,536. Positions: decoupled rotary dimensions, 64 per head. Cache: 576 numbers per token per layer, as much as grouped-query attention with 2.25 groups. Vocabulary: 102,400 tokens. Good at: long conversations served to many users at once. Not for: a single GPU

The family of the Transformer as this course covers it, with lesson 8, the KV cache and latent attention, highlighted and lessons 1 to 7 done The family of the Transformer as this course covers it, with lesson 8, the KV cache and latent attention, highlighted and lessons 1 to 7 done

References

  1. DeepSeek-AI (2024). DeepSeek-V2: a strong, economical, and efficient mixture-of-experts language model. arXiv:2405.04434. ↩
  2. Ji, T., Guo, B., Wu, Y., et al. (2025). Towards economical inference: enabling DeepSeek’s multi-head latent attention in any Transformer-based LLMs. arXiv:2502.14837. ↩
  3. Meng, F., Tang, P., Yao, Z., et al. (2025). TransMLA: multi-head latent attention is all you need. arXiv:2502.07864. ↩
  4. Liu, Z., Yuan, J., Jin, H., et al. (2024). KIVI: a tuning-free asymmetric 2bit quantization for KV cache. ICML 2024. ↩
  5. Kwon, W., Li, Z., Zhuang, S., et al. (2023). Efficient memory management for large language model serving with PagedAttention. SOSP 2023. ↩
  6. Character.AI, 2024. web.archive.org. ↩
  7. Brandon, W., Mishra, M., Nrusimha, A., et al. (2024). Reducing Transformer key-value cache size with cross-layer attention. arXiv:2405.12981. ↩
  8. Sun, Y., Dong, L., Zhu, Y., et al. (2024). You only cache once: decoder-decoder architectures for language models. NeurIPS 2024. ↩
  9. Zhang, Z., Sheng, Y., Zhou, T., et al. (2023). H2O: heavy-hitter oracle for efficient generative inference of large language models. NeurIPS 2023. ↩
  10. DeepSeek-AI (2024). DeepSeek-V3 technical report. arXiv:2412.19437.
  11. Kimi Team (2025). Kimi K2: open agentic intelligence. arXiv:2507.20534.
  12. Shazeer, N. (2019). Fast Transformer decoding: one write-head is all you need. arXiv:1911.02150.
  13. Ainslie, J., Lee-Thorp, J., de Jong, M., et al. (2023). GQA: training generalized multi-query Transformer models from multi-head checkpoints. EMNLP 2023.
  14. Hooper, C., Kim, S., Mohammadzadeh, H., et al. (2024). KVQuant: towards 10 million context length LLM inference with KV cache quantization. NeurIPS 2024.
  15. OpenBMB (2024). MiniCPM3-4B, model card and weights under the MiniCPM model licence.
]]> Eva Koroleva Mixture-of-experts LLMs from scratch with OLMoE 2026-08-24T16:00:00+00:00 2026-11-08T00:00:00+00:00 https://blog.korolevalab.com/blog/2026/transformer-architectures-07-moe Transformer architectures: from BERT to diffusion · lesson 7 of 15 · all lessons · run on Kaggle · open in Colab · code on GitHub

OLMoE-1B-7B has seven billion parameters and uses about one and a third billion of them for each token. Every one of its sixteen layers holds 64 feed-forward networks, the experts, and a small linear layer, the router, picks eight of them for each token.

Nobody tells the router what an expert is for. Whatever division of labour the experts end up with, they found it themselves during training on five trillion tokens.

In a dense model every token goes through every weight, so there is nothing to watch. A mixture can be watched: feed it Python and see which experts light up, feed it Russian, a language it barely saw in training, and watch again.

Then you can act on what you saw, throw away the experts a job doesn’t use, and keep a smaller model that still does the job.

A post on this blog built a mixture of experts from scratch, trained it next to dense models and compared the ways of keeping its experts busy. This lesson takes the next step, to a real model trained at scale: we read how nine production mixtures are built, write OLMoE in PyTorch, load the official weights and look inside.

A mixture-of-experts layer. A token's vector goes to the router, a linear layer that scores all 64 experts; the softmax turns the scores into probabilities, and the eight largest pick the experts that run. Each chosen expert is a small SwiGLU network; their outputs are added, each weighted by its probability, and the result joins the residual stream. The other 56 experts do nothing for this token A mixture-of-experts layer. A token's vector goes to the router, a linear layer that scores all 64 experts; the softmax turns the scores into probabilities, and the eight largest pick the experts that run. Each chosen expert is a small SwiGLU network; their outputs are added, each weighted by its probability, and the result joins the residual stream. The other 56 experts do nothing for this token

What you’ll do

  1. Read the configs of nine mixtures, from Mixtral to gpt-oss, and see how the design moved.
  2. Write the router four ways, find out which of them are the same function, and dispatch tokens to experts without computing what nobody asked for.
  3. Write OLMoE in PyTorch, load its official weights across two GPUs and match Hugging Face’s outputs.
  4. Route English, Python, school maths and Russian through it and find out which experts each one uses.
  5. Prune the experts a single kind of text doesn’t use, and measure what a smaller model costs.
  6. Measure how often a token needs the same experts as the token before it, which decides whether a mixture can live partly outside GPU memory.

What you need: two GPUs. On Kaggle, Settings → Accelerator → GPU T4 x2; OLMoE’s weights take 13 GB in 16-bit floats, more than one T4 has to spare. The notebook downloads OLMoE (14 GB), a slice of Russian Wikipedia (250 MB) and two small datasets; on two T4s it runs in about five minutes.

Setup

One small helper package, tfs, draws the course’s plots and token chips and knows where the notebook runs (Kaggle, Colab or your own machine). Everything the lesson teaches is written out in the notebook itself.

import importlib.util
import subprocess
import sys

TFS = "git+https://github.com/xmarva/transformers-from-scratch@main"
if importlib.util.find_spec("tfs") is None:
    subprocess.check_call([sys.executable, "-m", "pip", "install", "-q", TFS])

from tfs import env
from tfs.viz import table, token_chips, use_course_style

env.quiet()  # no progress bars or environment warnings in the outputs
SEED = 42
env.seed_everything(SEED)
COLORS = use_course_style()
WORK_DIR = env.work_dir()

import json
import re
import subprocess
import sys
import sysconfig
import time
from collections import OrderedDict
from html import escape
from pathlib import Path

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import pyarrow.parquet as pq
import torch
import torch.nn as nn
import torch.nn.functional as F
from IPython.display import HTML, display
from matplotlib.colors import LinearSegmentedColormap, to_rgb
from huggingface_hub import hf_hub_download
from safetensors.torch import load_file
from transformers import AutoTokenizer

SMOKE = env.smoke()  # CI runs the lesson on a small slice; you run it in full
device = "cuda" if torch.cuda.is_available() else "cpu"
GPUS = [f"cuda:{i}" for i in range(torch.cuda.device_count())] or ["cpu"]
HALF = torch.float16 if device == "cuda" else torch.float32  # 16-bit matmuls crawl on a CPU
assert device == "cuda" or SMOKE, "This lesson runs a 7B model: use two GPUs"
for name in GPUS:
    if name != "cpu":
        print(name, torch.cuda.get_device_name(name))

cuda:0 Tesla T4
cuda:1 Tesla T4

1. Nine mixtures side by side

As in lesson 6, the fastest way to read an architecture is its config.json. Here are the configs of nine mixtures released between December 2023 and August 2025, with the total and active parameter counts their papers and model cards give. Every model calls its settings a little differently, so the reading takes a few lines.

# repository, name, released, total / active parameters from the paper or the model card
MOE_MODELS = [
    ("mistralai/Mixtral-8x7B-v0.1", "Mixtral 8x7B", "2023-12", "47B / 13B"),
    ("Qwen/Qwen1.5-MoE-A2.7B", "Qwen1.5-MoE", "2024-03", "14.3B / 2.7B"),
    ("deepseek-ai/DeepSeek-V2-Lite", "DeepSeek-V2-Lite", "2024-05", "15.7B / 2.4B"),
    ("allenai/OLMoE-1B-7B-0924", "OLMoE", "2024-09", "6.9B / 1.3B"),
    ("ibm-granite/granite-3.1-3b-a800m-base", "Granite 3.1 3B", "2024-12", "3.3B / 0.8B"),
    ("deepseek-ai/DeepSeek-V3", "DeepSeek-V3", "2024-12", "671B / 37B"),
    ("Qwen/Qwen3-30B-A3B", "Qwen3-30B-A3B", "2025-04", "30.5B / 3.3B"),
    ("openai/gpt-oss-20b", "gpt-oss-20b", "2025-08", "20.9B / 3.6B"),
    ("openai/gpt-oss-120b", "gpt-oss-120b", "2025-08", "116.8B / 5.1B"),
]


def read_config(repo):
    with open(hf_hub_download(repo, "config.json")) as f:
        return json.load(f)


def first(c, *keys, default=0):
    for key in keys:
        if c.get(key) is not None:
            return c[key]
    return default


def router_order(c):
    kind = c["model_type"]
    if kind in ("granitemoe", "gpt_oss"):
        return "top-k, softmax"
    if c.get("scoring_func") == "sigmoid":
        return "sigmoid + bias, top-k, renormalise"
    renorm = kind == "mixtral" or c.get("norm_topk_prob")
    return "softmax, top-k" + (", renormalise" if renorm else "")


configs = {name: read_config(repo) for repo, name, _, _ in MOE_MODELS}
rows = []
for _, name, released, sizes in MOE_MODELS:
    c = configs[name]
    experts = first(c, "num_local_experts", "num_experts", "n_routed_experts")
    top_k = c["num_experts_per_tok"]
    shared = first(c, "n_shared_experts")
    if c.get("shared_expert_intermediate_size"):
        shared = c["shared_expert_intermediate_size"] // first(c, "moe_intermediate_size")
    rows.append(
        {
            "model": name,
            "released": released,
            "total / active": sizes,
            "layers": c["num_hidden_layers"],
            "dense layers": first(c, "first_k_dense_replace"),
            "experts": experts,
            "per token": top_k,
            "shared": shared,
            "expert width": first(c, "moe_intermediate_size", "intermediate_size"),
            "model width": c["hidden_size"],
            "router": router_order(c),
        }
    )
table(pd.DataFrame(rows))
modelreleasedtotal / activelayersdense layersexpertsper tokensharedexpert widthmodel widthrouter
Mixtral 8x7B2023-1247B / 13B32082014,3364,096softmax, top-k, renormalise
Qwen1.5-MoE2024-0314.3B / 2.7B24060441,4082,048softmax, top-k
DeepSeek-V2-Lite2024-0515.7B / 2.4B27164621,4082,048softmax, top-k
OLMoE2024-096.9B / 1.3B16064801,0242,048softmax, top-k
Granite 3.1 3B2024-123.3B / 0.8B32040805121,536top-k, softmax
DeepSeek-V32024-12671B / 37B613256812,0487,168sigmoid + bias, top-k, renormalise
Qwen3-30B-A3B2025-0430.5B / 3.3B480128807682,048softmax, top-k, renormalise
gpt-oss-20b2025-0820.9B / 3.6B24032402,8802,880top-k, softmax
gpt-oss-120b2025-08116.8B / 5.1B360128402,8802,880top-k, softmax

Three things changed in under two years.

The experts got smaller and more numerous. Mixtral has eight experts as wide as a dense model’s feed-forward network and sends each token to two. OLMoE, Qwen3 and DeepSeek-V3 use 64 to 256 narrow experts and send each token to eight.

DeepSeek’s argument for this, fine-grained experts, is combinatorial: two experts out of sixteen can be chosen in 120 ways, eight out of 64 in “4,426,165,368”[1], so a token gets a mix much closer to what it needs. gpt-oss went the other way, to 32 or 128 wide experts and four per token.

Some models keep experts that every token uses. DeepSeek’s shared experts hold the knowledge every token needs, so that the routed ones don’t each keep a copy of it: “Regardless of the router module, each token will be deterministically assigned to these shared experts.”

Qwen1.5-MoE had them; Qwen3 dropped them; OLMoE tested one and found that “sharing an expert performs slightly worse”[2].

The routers disagree on the order of operations. Some take a softmax over all experts and keep the top \(k\) probabilities; some keep the top \(k\) scores and take a softmax over those; DeepSeek-V3 scores experts with a sigmoid and adds a bias that only affects which experts are chosen. Section 2 shows which of these differences matter.

Under the routing, every model here is the decoder of lesson 6 with its feed-forward network replaced. Attention, norms and positions don’t change, which is why the table has no column for them.

2. The router and the experts

A mixture-of-experts layer computes, for a token vector \(x\) and experts \(E_1, \dots, E_n\),

\[y = \sum_{i \in \mathcal{T}} g_i(x)\, E_i(x),\]

where \(\mathcal{T}\) is the set of the \(k\) chosen experts and \(g_i\) their weights. The router is a single matrix \(W_r\) with one row per expert; its scores are \(s = W_r x\). The four routers of the table differ only in how they turn \(s\) into \(\mathcal{T}\) and \(g\).

def softmax_topk(s, k):  # OLMoE, Qwen1.5-MoE, DeepSeek-V2-Lite
    weights, chosen = s.softmax(dim=-1).topk(k, dim=-1)
    return weights, chosen


def softmax_topk_renorm(s, k):  # Mixtral, Qwen3
    weights, chosen = softmax_topk(s, k)
    return weights / weights.sum(dim=-1, keepdim=True), chosen


def topk_softmax(s, k):  # Granite, gpt-oss
    top, chosen = s.topk(k, dim=-1)
    return top.softmax(dim=-1), chosen


def sigmoid_bias(s, k, bias):  # DeepSeek-V3: the bias picks the experts, not their weights
    affinity = s.sigmoid()
    chosen = (affinity + bias).topk(k, dim=-1).indices
    weights = affinity.gather(-1, chosen)
    return weights / weights.sum(dim=-1, keepdim=True), chosen


torch.manual_seed(0)
s = torch.randn(1000, 64) * 2  # scores of 1,000 tokens for 64 experts
results = {
    "softmax, top-k": softmax_topk(s, 8),
    "softmax, top-k, renormalise": softmax_topk_renorm(s, 8),
    "top-k, softmax": topk_softmax(s, 8),
    "sigmoid + bias, top-k, renormalise": sigmoid_bias(s, 8, torch.zeros(64)),
}
reference_w, reference_i = results["top-k, softmax"]
rows = []
for name, (w, i) in results.items():
    rows.append(
        {
            "router": name,
            "same experts as top-k, softmax": bool(torch.equal(i, reference_i)),
            "largest weight difference": f"{(w - reference_w).abs().max().item():.2e}",
            "weights sum to": f"{w.sum(dim=-1).mean().item():.3f}",
        }
    )
table(pd.DataFrame(rows))
routersame experts as top-k, softmaxlargest weight differenceweights sum to
softmax, top-kTrue2.23e-010.744
softmax, top-k, renormaliseTrue1.79e-071.000
top-k, softmaxTrue0.00e+001.000
sigmoid + bias, top-k, renormaliseTrue8.51e-011.000

All four routers choose the same experts here, and that is no accident: a softmax and a sigmoid never change the order of the scores, and without a bias the top \(k\) scores stay the top \(k\) under any increasing function. They differ in the weights.

Two of them are in fact one function: a softmax over all experts followed by renormalising the top \(k\) gives exactly what a softmax over the top \(k\) scores gives, because the experts left out cancel between numerator and denominator.

Mixtral and Qwen3 on one side, Granite and gpt-oss on the other, write the same router in two orders. OLMoE’s is the odd one out: its eight weights add up to less than one, by however much probability the router gave to the experts it didn’t choose.

DeepSeek-V3’s sigmoid gives yet other weights, and its bias, which is zero here, is what lets it steer the choice without touching them.

The experts themselves are lesson 5’s SwiGLU network of the first course, one per expert, each much narrower than a dense model’s. The naive way to run a layer is to run every expert on every token and multiply most of the results by zero.

The way that saves the compute groups the tokens by expert: sort the (token, expert) pairs by expert, run each expert once on its own group, and add each result back to its token, weighted.

class Experts(nn.Module):
    """n SwiGLU networks kept as two stacked weights, one row of experts per tensor."""

    def __init__(self, n, width, hidden):
        super().__init__()
        self.gate_up = nn.Parameter(
            torch.empty(n, 2 * hidden, width)
        )  # gate rows, then up rows
        self.down = nn.Parameter(torch.empty(n, width, hidden))

    def run(self, e, x):
        gate, up = (x @ self.gate_up[e].T).chunk(2, dim=-1)
        return (F.silu(gate) * up) @ self.down[e].T


def dispatch(experts, x, weights, chosen):
    """Run each expert once, on the tokens that chose it; add the results back, weighted."""
    k = chosen.shape[-1]
    pairs = chosen.flatten()
    order = pairs.argsort(stable=True)
    tokens = order // k  # the token each sorted pair came from
    w = weights.flatten()[order]
    counts = pairs.bincount(minlength=experts.gate_up.shape[0]).tolist()
    out = torch.zeros_like(x)
    start = 0
    for e, n in enumerate(counts):
        if n:
            rows = tokens[start : start + n]
            out.index_add_(
                0, rows, (experts.run(e, x[rows]) * w[start : start + n, None]).to(x.dtype)
            )
        start += n
    return out


def dense_run(experts, x, weights, chosen):
    """The same layer computed the wasteful way: every expert on every token."""
    n = experts.gate_up.shape[0]
    gate = torch.zeros(x.shape[0], n, dtype=weights.dtype).scatter(1, chosen, weights)
    outputs = torch.stack(
        [experts.run(e, x) for e in range(n)], dim=1
    )  # tokens x experts x width
    return (outputs * gate[..., None]).sum(dim=1)


torch.manual_seed(0)
toy = Experts(64, 256, 128)
nn.init.normal_(toy.gate_up, std=0.05)
nn.init.normal_(toy.down, std=0.05)
x = torch.randn(512, 256)
weights, chosen = softmax_topk(x @ torch.randn(256, 64) * 0.1, 8)
timings = {}
for name, fn in [
    ("every expert on every token", dense_run),
    ("grouped by expert", dispatch),
]:
    start = time.time()
    with torch.no_grad():
        y = fn(toy, x, weights, chosen)
    timings[name] = (time.time() - start, y)
difference = (
    (timings["grouped by expert"][1] - timings["every expert on every token"][1])
    .abs()
    .max()
)
for name, (seconds, _) in timings.items():
    print(f"{name:28} {seconds * 1000:7.1f} ms")
print("max difference:", f"{difference.item():.2e}")

every expert on every token    153.9 ms
grouped by expert               63.0 ms
max difference: 5.96e-08

The two ways agree to the last digits a 32-bit float carries, and grouping by expert is already faster on a CPU even at this toy size, though each expert sees only a few dozen tokens.

The gap grows with the number of experts: the wasteful way pays for all of them, the grouped one for the eight a token chose.

Production kernels go further and run all the groups in one grouped matrix multiplication, but the arithmetic is the one above. How the router is kept from sending everything to a few favourites, the auxiliary loss of Switch Transformer and the bias of DeepSeek-V3, is in the blog post linked at the top; OLMoE was trained with the loss.

3. OLMoE in PyTorch

OLMoE comes from the Allen Institute for AI, which published everything needed to rebuild it: the weights, the training data, the code and the logs. It is “a fully open, state-of-the-art language model” with “6.9B total parameters, of which only 1.3B are activated”, trained on “5.1 trillion tokens”[2]. Its config:

REPO = "allenai/OLMoE-1B-7B-0924"
olmoe_config = configs["OLMoE"]
olmoe_config.setdefault("rms_norm_eps", 1e-5)  # not in the config: the library default
keys = [
    "hidden_size",
    "num_hidden_layers",
    "num_attention_heads",
    "num_key_value_heads",
    "num_experts",
    "num_experts_per_tok",
    "intermediate_size",
    "vocab_size",
    "norm_topk_prob",
    "tie_word_embeddings",
]
for key in keys:
    print(f"{key:22} {olmoe_config[key]}")
print(f"{'rope_theta':22} {olmoe_config.get('rope_theta')}")

hidden_size            2048
num_hidden_layers      16
num_attention_heads    16
num_key_value_heads    16
num_experts            64
num_experts_per_tok    8
intermediate_size      1024
vocab_size             50304
norm_topk_prob         False
tie_word_embeddings    False
rope_theta             10000.0

Sixteen heads with sixteen key-value heads: OLMoE uses plain multi-head attention, not lesson 6’s grouped queries. It does use QK-norm, with one difference from Qwen3: the norm runs over the whole query and key projections, 2,048 numbers each, before they are split into heads.

The router takes a softmax over all 64 experts and keeps the eight largest probabilities without renormalising them, so a token whose router is unsure gets a smaller total update from its experts than one whose router is sure.

Everything outside the experts is lesson 6’s code with those two changes. One more change is for the GPU, not the model: the residual stream is kept in 32-bit floats between layers.

OLMoE was trained in bfloat16, which has a much wider range than the float16 a T4 offers, and the residual stream, a running sum over 32 sublayers, is where large values build up.

class RMSNorm(nn.Module):
    def __init__(self, dim, eps):
        super().__init__()
        self.eps = eps
        self.weight = nn.Parameter(torch.ones(dim))

    def forward(self, x):
        x32 = x.float()
        x32 = x32 * torch.rsqrt(x32.pow(2).mean(dim=-1, keepdim=True) + self.eps)
        return self.weight * x32.to(self.weight.dtype)


def rope_tables(positions, head_dim, theta):
    exponents = torch.arange(0, head_dim, 2, device=positions.device).float() / head_dim
    angles = positions.float()[..., None] / theta**exponents
    angles = torch.cat([angles, angles], dim=-1)
    return angles.cos()[:, None], angles.sin()[:, None]


def rotate(x, cos, sin):
    half = x.shape[-1] // 2
    rotated = torch.cat([-x[..., half:], x[..., :half]], dim=-1)
    return x * cos.to(x.dtype) + rotated * sin.to(x.dtype)


class Attention(nn.Module):
    def __init__(self, c):
        super().__init__()
        d = c["hidden_size"]
        self.heads = c["num_attention_heads"]
        self.head_dim = d // self.heads
        self.q_proj = nn.Linear(d, d, bias=False)
        self.k_proj = nn.Linear(d, d, bias=False)
        self.v_proj = nn.Linear(d, d, bias=False)
        self.o_proj = nn.Linear(d, d, bias=False)
        self.q_norm = RMSNorm(
            d, c["rms_norm_eps"]
        )  # over the whole projection, before the heads
        self.k_norm = RMSNorm(d, c["rms_norm_eps"])

    def forward(self, x, cos, sin, mask):
        b, t, _ = x.shape
        shape = (b, t, self.heads, self.head_dim)
        q = rotate(self.q_norm(self.q_proj(x)).view(shape).transpose(1, 2), cos, sin)
        k = rotate(self.k_norm(self.k_proj(x)).view(shape).transpose(1, 2), cos, sin)
        v = self.v_proj(x).view(shape).transpose(1, 2)
        out = F.scaled_dot_product_attention(q, k, v, attn_mask=mask)
        return self.o_proj(out.transpose(1, 2).reshape(b, t, -1))


class MoE(nn.Module):
    def __init__(self, c):
        super().__init__()
        d, n = c["hidden_size"], c["num_experts"]
        self.k = c["num_experts_per_tok"]
        self.gate = nn.Linear(d, n, bias=False)  # the router
        self.experts = Experts(n, d, c["intermediate_size"])
        self.keep = None  # section 5: a mask of the experts a pruned model still has
        self.trace = None  # section 4: a list that collects every token's chosen experts

    def forward(self, x):
        b, t, d = x.shape
        flat = x.reshape(-1, d)
        scores = self.gate(flat).float()
        if self.keep is not None:
            scores = scores.masked_fill(~self.keep, float("-inf"))
        weights, chosen = softmax_topk(scores, self.k)  # no renormalisation in OLMoE
        if self.trace is not None:
            self.trace.append(chosen.to(torch.uint8).cpu())
        return dispatch(self.experts, flat, weights, chosen).view(b, t, d)


class Block(nn.Module):
    def __init__(self, c):
        super().__init__()
        self.input_layernorm = RMSNorm(c["hidden_size"], c["rms_norm_eps"])
        self.self_attn = Attention(c)
        self.post_attention_layernorm = RMSNorm(c["hidden_size"], c["rms_norm_eps"])
        self.mlp = MoE(c)

    def forward(self, x, cos, sin, mask):  # x: the residual stream, in 32-bit floats
        x = x + self.self_attn(self.input_layernorm(x), cos, sin, mask).float()
        return x + self.mlp(self.post_attention_layernorm(x)).float()


class OLMoE(nn.Module):
    def __init__(self, c):
        super().__init__()
        self.config = c
        self.embed_tokens = nn.Embedding(c["vocab_size"], c["hidden_size"])
        self.layers = nn.ModuleList([Block(c) for _ in range(c["num_hidden_layers"])])
        self.norm = RMSNorm(c["hidden_size"], c["rms_norm_eps"])
        self.lm_head = nn.Linear(c["hidden_size"], c["vocab_size"], bias=False)

    def hidden(self, ids, mask):
        c = self.config
        positions = torch.arange(ids.shape[1], device=ids.device)[None]
        x = self.embed_tokens(ids).float()
        tables = {}
        for layer in self.layers:
            where = layer.input_layernorm.weight.device  # layers may sit on different GPUs
            if where not in tables:
                theta = c.get("rope_theta") or 10000.0
                tables[where] = rope_tables(
                    positions.to(where), c["hidden_size"] // c["num_attention_heads"], theta
                )
            x = layer(x.to(where), *tables[where], mask.to(where))
        return self.norm(x.to(self.norm.weight.device))

    def logits(self, h):
        return self.lm_head(h.to(self.lm_head.weight.device)).float()

    def forward(self, ids, mask):
        return self.logits(self.hidden(ids, mask))

The checkpoint stores each expert’s three matrices separately, 3,072 tensors in all; we stack them into the two tensors per layer that Experts keeps.

The weights go straight from each file to the GPU that will use them, so the 13 GB never sit in the machine’s memory at once: the first half of the layers lives on the first T4, the second half and the output layer on the second.

def placement(c, gpus):
    n = c["num_hidden_layers"]
    return {i: gpus[i * len(gpus) // n] for i in range(n)}


def checkpoint_files(source):
    if Path(source).is_dir():  # a model saved to disk (the smoke run builds a small one)
        return sorted(str(p) for p in Path(source).glob("*.safetensors"))
    try:
        with open(hf_hub_download(source, "model.safetensors.index.json")) as f:
            names = sorted(set(json.load(f)["weight_map"].values()))
    except Exception:
        names = ["model.safetensors"]
    return [hf_hub_download(source, name) for name in names]


def load_olmoe(source, c, dtype, gpus):
    where = placement(c, gpus)
    with torch.device("meta"):
        model = OLMoE(c)

    def target(key):
        match = re.match(r"layers\.(\d+)\.", key)
        return (
            where[int(match.group(1))]
            if match
            else (gpus[0] if "embed" in key else gpus[-1])
        )

    state, pending = {}, {}
    n = c["num_experts"]
    for path in checkpoint_files(source):
        for key, value in load_file(path).items():
            key = key.removeprefix("model.")
            expert = re.match(
                r"(layers\.\d+\.mlp\.experts)\.(\d+)\.(gate|up|down)_proj\.weight", key
            )
            if expert:  # one expert's matrix: wait for all of them, then stack
                prefix, e, part = expert.group(1), int(expert.group(2)), expert.group(3)
                pending.setdefault(prefix, {})[(e, part)] = value
                parts = pending[prefix]
                if len(parts) == 3 * n:
                    device = target(prefix)
                    gate_up = torch.stack(
                        [
                            torch.cat([parts[(e, "gate")], parts[(e, "up")]])
                            for e in range(n)
                        ]
                    )
                    down = torch.stack([parts[(e, "down")] for e in range(n)])
                    state[prefix + ".gate_up"] = gate_up.to(device, dtype)
                    state[prefix + ".down"] = down.to(device, dtype)
                    del pending[prefix]
                continue
            key = key.replace("experts.gate_up_proj", "experts.gate_up")  # already stacked
            key = key.replace("experts.down_proj", "experts.down")
            state[key] = value.to(target(key), dtype)
    model.load_state_dict(state, strict=True, assign=True)
    return model.eval()

The check, as always: the same tokens through our model and through Hugging Face’s. Two copies of OLMoE don’t fit on two T4s together, so Hugging Face’s runs first, in a separate process, saves its logits and exits.

tok = AutoTokenizer.from_pretrained(REPO)
text = (
    "The Rhine rises in the Swiss Alps, flows through Lake Constance and along "
    "the French–German border, and reaches the North Sea through a delta in the "
    "Netherlands."
)
check_ids = tok(text, return_tensors="pt").input_ids
torch.save(check_ids, WORK_DIR / "check_ids.pt")

if SMOKE:  # a small OLMoE with random weights stands in for the real one
    from transformers import OlmoeConfig, OlmoeForCausalLM

    small = OlmoeConfig(
        vocab_size=olmoe_config["vocab_size"],
        hidden_size=64,
        intermediate_size=32,
        num_hidden_layers=2,
        num_attention_heads=4,
        num_key_value_heads=4,
        num_experts=16,
        num_experts_per_tok=4,
        eos_token_id=0,
        pad_token_id=1,
    )
    torch.manual_seed(0)
    reference = OlmoeForCausalLM(small).eval()
    SOURCE = WORK_DIR / "olmoe-small"
    reference.save_pretrained(SOURCE)
    with torch.no_grad():
        theirs = reference(check_ids).logits.float()
    olmoe_config = json.loads((SOURCE / "config.json").read_text())
    olmoe_config.setdefault("rope_theta", 10000.0)
else:
    SOURCE = REPO
    script = f"""
import torch
from transformers import AutoModelForCausalLM
model = AutoModelForCausalLM.from_pretrained(
    "{REPO}", dtype=torch.float16, device_map="auto")
ids = torch.load("{WORK_DIR / "check_ids.pt"}").to("cuda:0")
with torch.no_grad():
    torch.save(model(ids).logits.float().cpu(), "{WORK_DIR / "theirs.pt"}")
"""
    subprocess.run([sys.executable, "-c", script], check=True)
    theirs = torch.load(WORK_DIR / "theirs.pt")

model = load_olmoe(SOURCE, olmoe_config, HALF, GPUS)
t = check_ids.shape[1]
causal = torch.ones(t, t, dtype=torch.bool).tril()
with torch.no_grad():
    ours = model(check_ids.to(GPUS[0]), causal).cpu()
print("tokens:", t)
print("largest logit:", f"{theirs.abs().max().item():.2f}")
print("max difference:", f"{(ours - theirs).abs().max().item():.2e}")
print(
    "same next-token prediction everywhere:",
    bool((ours.argmax(-1) == theirs.argmax(-1)).all()),
)

tokens: 35
largest logit: 23.70
max difference: 6.41e-01
same next-token prediction everywhere: True

Every next-token prediction is the same as Hugging Face’s. The logits differ more than in earlier lessons, for two reasons that have nothing to do with the arithmetic being wrong; the smoke test of this notebook, which runs a small random OLMoE in 32-bit floats, agrees with Hugging Face’s to the seventh digit.

Here both models run in 16-bit floats, ours with the residual stream in 32 bits and theirs without, so their rounding differs. And a mixture amplifies small differences: when two experts have nearly the same score for the eighth place, a rounding error decides which one runs, and the token gets a different expert’s output.

Where do the parameters live?

parts = {"embeddings and output": 0, "attention": 0, "routers": 0, "experts": 0, "norms": 0}
for name, p in model.named_parameters():
    if "experts" in name:
        parts["experts"] += p.numel()
    elif "mlp.gate" in name:
        parts["routers"] += p.numel()
    elif "self_attn" in name and "norm" not in name:
        parts["attention"] += p.numel()
    elif "embed" in name or "lm_head" in name:
        parts["embeddings and output"] += p.numel()
    else:
        parts["norms"] += p.numel()
total = sum(parts.values())
c = olmoe_config
per_expert = 3 * c["hidden_size"] * c["intermediate_size"]
active = (
    total
    - parts["experts"]
    + c["num_hidden_layers"] * c["num_experts_per_tok"] * per_expert
)
table(
    pd.DataFrame(
        {
            "part": list(parts),
            "parameters": list(parts.values()),
            "share": [f"{n / total:.1%}" for n in parts.values()],
        }
    )
)
print(f"used for one token: {active:,} of {total:,} ({active / total:.0%})")

used for one token: 1,282,017,280 of 6,919,161,856 (19%)

Nearly all of OLMoE is experts. A token runs through all the attention, the embeddings and one eighth of the experts, a fifth of the model in all, which is where the 1B in its name comes from. The routers that decide all of this are a rounding error: one matrix of 64 rows per layer.

4. Which experts fire

Four kinds of text, each from its own corpus. English is Wikipedia articles from WikiText-103, as in lesson 6. Python is the source code of Python’s own standard library, which is on every machine that runs this notebook.

Maths is the 1,319 school word problems of GSM8K’s test set with their worked solutions[3], which mix English, numbers and arithmetic. Russian is articles from a slice of Russian Wikipedia.

OLMoE’s training data is overwhelmingly English, and its tokenizer shows it: a Russian word falls apart into many more tokens than an English one. Each text is cut into windows of 512 tokens; most go to measuring how the experts are used, and a few are held out for section 5.

LENGTH = 64 if SMOKE else 512
CALIBRATION, HELD_OUT = (2, 1) if SMOKE else (48, 16)


def windows(texts, n):
    ids = []
    for text in texts:
        ids.extend(tok(text, add_special_tokens=False).input_ids + [tok.eos_token_id])
        if len(ids) >= n * LENGTH:
            break
    ids = ids[: n * LENGTH]
    return torch.tensor(ids).view(n, LENGTH)


def english_texts():
    name = "wikitext-103-raw-v1/test-00000-of-00001.parquet"
    path = hf_hub_download("Salesforce/wikitext", name, repo_type="dataset")
    return ["".join(pd.read_parquet(path)["text"].tolist())]


def python_texts():
    root = Path(sysconfig.get_paths()["stdlib"])
    modules = [
        "argparse",
        "difflib",
        "fractions",
        "statistics",
        "textwrap",
        "heapq",
        "dataclasses",
        "pathlib",
        "calendar",
        "ipaddress",
        "csv",
        "shlex",
        "string",
        "gettext",
        "tokenize",
    ]
    files = [root / f"{m}.py" for m in modules if (root / f"{m}.py").exists()]
    files += [p for p in sorted(root.glob("json/*.py")) + sorted(root.glob("email/*.py"))]
    return [f.read_text(errors="ignore") for f in files]


def maths_texts():
    path = hf_hub_download(
        "openai/gsm8k", "main/test-00000-of-00001.parquet", repo_type="dataset"
    )
    frame = pd.read_parquet(path)
    # the solutions carry calculator annotations, <<48/2=24>>, that no reader sees
    return [
        q + "\n" + re.sub(r"<<[^>]*>>", "", a) + "\n"
        for q, a in zip(frame.question, frame.answer)
    ]


def russian_texts():
    path = hf_hub_download(
        "wikimedia/wikipedia",
        "20231101.ru/train-00002-of-00021.parquet",
        repo_type="dataset",
    )
    texts = (
        pq.ParquetFile(path).read_row_group(0, columns=["text"]).column("text").to_pylist()
    )
    return [t for t in texts if len(t) > 2000]  # whole articles, not stubs


corpora = {
    "English": english_texts(),
    "Python": python_texts(),
    "maths": maths_texts(),
    "Russian": russian_texts(),
}
DOMAINS = list(corpora)
data = {}
for name, texts in corpora.items():
    all_ids = windows(texts, CALIBRATION + HELD_OUT)
    data[name] = {"calibration": all_ids[:CALIBRATION], "held out": all_ids[CALIBRATION:]}
words = {
    "English": "The river flows through the old town and into the sea.",
    "Russian": "Река течёт через старый город и впадает в море.",
}
for language, sentence in words.items():
    print(
        f"{language:8} {len(sentence.split())} words -> "
        f"{len(tok(sentence).input_ids)} tokens"
    )
print({name: tuple(d["calibration"].shape) for name, d in data.items()})

English  11 words -> 12 tokens
Russian  9 words -> 26 tokens
{'English': (48, 512), 'Python': (48, 512), 'maths': (48, 512), 'Russian': (48, 512)}

Now the routing. A trace is switched on in every layer’s router, each corpus goes through the model, and every token leaves behind the eight experts it chose in each layer.

BATCH = 2 if SMOKE else 8
n_layers, n_experts = olmoe_config["num_hidden_layers"], olmoe_config["num_experts"]


def causal_mask(ids):
    t = ids.shape[1]
    return torch.ones(t, t, dtype=torch.bool).tril()


def route(batches):
    """Every token's chosen experts in every layer: tokens x layers x k."""
    for layer in model.layers:
        layer.mlp.trace = []
    with torch.no_grad():
        for i in range(0, len(batches), BATCH):
            chunk = batches[i : i + BATCH]
            model.hidden(chunk.to(GPUS[0]), causal_mask(chunk))
    traces = torch.stack([torch.cat(layer.mlp.trace) for layer in model.layers], dim=1)
    for layer in model.layers:
        layer.mlp.trace = None
    return traces


start = time.time()
traces = {name: route(d["calibration"]) for name, d in data.items()}
print(
    f"routed {sum(len(t) for t in traces.values()):,} tokens in {time.time() - start:.0f} s"
)

# usage[d, l, e]: the share of domain d's tokens that chose expert e in layer l
usage = np.zeros((len(DOMAINS), n_layers, n_experts))
for d, name in enumerate(DOMAINS):
    t = traces[name].long().numpy()
    for layer in range(n_layers):
        usage[d, layer] = np.bincount(t[:, layer].flatten(), minlength=n_experts) / len(t)

routed 98,304 tokens in 18 s

Each token picks eight experts, so a domain’s usage in a layer adds up to eight.

To ask what an expert prefers, we compare its usage across the four domains: an expert used equally by all four has no preference; one used by a single domain and nobody else is that domain’s expert. The heat map shows three layers, with the experts sorted by the domain they prefer.

preference = usage / usage.sum(axis=0, keepdims=True).clip(
    min=1e-9
)  # share per domain, per expert
shades = LinearSegmentedColormap.from_list("share", [COLORS["grid"], COLORS["att"]])
shown_layers = [0, n_layers // 2, n_layers - 1]
fig, axes = plt.subplots(
    len(shown_layers), 1, figsize=(10, 1.15 * len(shown_layers) * len(DOMAINS) / 2.4)
)
for ax, layer in zip(np.atleast_1d(axes), shown_layers):
    p = preference[:, layer]
    order = np.lexsort((-p.max(axis=0), p.argmax(axis=0)))
    ax.imshow(p[:, order], aspect="auto", cmap=shades, vmin=0, vmax=1)
    ax.set_yticks(range(len(DOMAINS)), DOMAINS)
    ax.set_xticks([])
    ax.set_title(
        f"layer {layer}: one column per expert, its use split by kind of text",
        fontsize=10,
        loc="left",
    )
plt.tight_layout()
plt.show()

Output of the code above

The experts have divided the work, and they did it in every layer. In each panel a block of experts belongs to each kind of text, dark in one row and pale in the others.

The blocks sharpen with depth: in the first layer many experts are still shared, while in the last most of them serve one kind of text and ignore the rest.

Russian has its own experts from the very first layer, which makes sense for a language whose tokens look nothing like the others: Cyrillic letters and fragments of words that English text never produces.

The same thing counted across all sixteen layers: in each layer, how many experts get more than half of their use from one kind of text.

threshold = 0.5
specialists = np.stack(
    [
        (preference.argmax(axis=0) == d) & (preference.max(axis=0) > threshold)
        for d in range(len(DOMAINS))
    ]
).sum(axis=2)  # domains x layers
fig, ax = plt.subplots(figsize=(9, 3.4))
bottom = np.zeros(n_layers)
for d, name in enumerate(DOMAINS):
    ax.bar(
        range(n_layers),
        specialists[d],
        bottom=bottom,
        label=name,
        color=[COLORS["enc"], COLORS["att"], COLORS["data"], COLORS["dec"]][d],
    )
    bottom += specialists[d]
ax.set_xlabel("layer")
ax.set_ylabel(f"experts of {n_experts}")
ax.set_title("Experts that get more than half of their use from one kind of text")
ax.legend(frameon=False, ncol=4, loc="upper center", bbox_to_anchor=(0.5, -0.2))
plt.show()

Experts that get more than half of their use from one kind of text

The count grows towards the last layers, where more than half of the experts take most of their work from a single kind of text.

This is what OLMoE’s authors found on their own data, “experts exhibit domain and vocabulary specialization”, and it is the opposite of what Mistral found in Mixtral, where “the selection of experts appears to be more aligned with the syntax rather than the domain”.

OLMoE’s authors have an explanation: Mixtral was not trained as a mixture from scratch but “upcycled from Mistral”, its experts copied from one dense network, and they never grew apart as much.

Specialisation also shows from the other side. How many of the 64 experts does a domain need in a layer to cover most of its routing? A domain that spreads evenly over all of them needs most of the layer; a domain with its own experts needs few.

def footprint(u, share):
    """How many experts, taken from the most used down, carry `share` of the routing."""
    ordered = np.sort(u)[::-1]
    return int(np.searchsorted(np.cumsum(ordered) / ordered.sum(), share) + 1)


rows = []
for d, name in enumerate(DOMAINS):
    for share in [0.5, 0.9]:
        counts = [footprint(usage[d, layer], share) for layer in range(n_layers)]
        rows.append(
            {
                "text": name,
                "share of routing": f"{share:.0%}",
                "experts needed, fewest layer": min(counts),
                "median layer": int(np.median(counts)),
                "most layer": max(counts),
            }
        )
table(pd.DataFrame(rows))
textshare of routingexperts needed, fewest layermedian layermost layer
English50%161921
English90%434753
Python50%6717
Python90%303349
maths50%111419
maths90%394451
Russian50%5915
Russian90%143745

English spreads over the most experts: covering half of its routing takes a quarter to a third of a layer, and covering nine tenths takes most of the layer.

Python and Russian are concentrated: a handful of experts carry half of their routing in a typical layer. That is the difference between a model’s main language, which uses the whole model, and the kinds of text it can afford to delegate to a few specialists.

Last, the tokens themselves. Below, a sentence of each kind with every token coloured by the domain its first expert prefers, in the first and the last layer.

samples = {
    "English": "The museum opens at nine and closes at six on weekdays.",
    "Python": "for line in open(path):\n    if line.startswith('#'):\n        continue",
    "maths": "She buys 3 boxes of 12 eggs and uses 7, so 36 - 7 = 29 are left.",
    "Russian": "Музей открывается в девять утра и закрывается в шесть вечера.",
}
palette = [COLORS["enc"], COLORS["att"], COLORS["data"], COLORS["dec"]]


def chip(piece, color):
    r, g, b = (int(255 * v) for v in to_rgb(color))
    return (
        f'<span style="background:rgba({r},{g},{b},0.32);border-radius:4px;padding:1px 3px;'
        f'margin:1px;font-family:monospace;white-space:pre">{escape(piece)}</span>'
    )


lines = ['<div style="line-height:2.1">']
for name, sentence in samples.items():
    ids = tok(sentence, return_tensors="pt").input_ids
    trace = route(ids)
    pieces = [tok.decode([i]).replace("\n", "↵") for i in ids[0].tolist()]
    lines.append(f"<div><b>{name}</b></div>")
    for layer in [0, n_layers - 1]:
        chips = []
        for piece, chosen_experts in zip(pieces, trace[:, layer]):
            liked = int(preference[:, layer, int(chosen_experts[0])].argmax())
            chips.append(chip(piece, palette[liked]))
        lines.append(
            f'<div><span style="opacity:0.6">layer {layer}&nbsp;</span>'
            f'{"".join(chips)}</div>'
        )
legend = "".join(chip(d, palette[i]) for i, d in enumerate(DOMAINS))
caption = '<span style="opacity:0.6">the first expert prefers&nbsp;</span>'
lines.append(f"<div>{caption}{legend}</div></div>")
display(HTML("".join(lines)))
English
layer 0 The museum opens at nine and closes at six on weekdays.
layer 15 The museum opens at nine and closes at six on weekdays.
Python
layer 0 for line in open(path):↵ if line.startswith('#'):↵ continue
layer 15 for line in open(path):↵ if line.startswith('#'):↵ continue
maths
layer 0 She buys 3 boxes of 12 eggs and uses 7, so 36 - 7 = 29 are left.
layer 15 She buys 3 boxes of 12 eggs and uses 7, so 36 - 7 = 29 are left.
Russian
layer 0 Музей открывается в девять утра и закрывается в шесть вечера.
layer 15 Музей открывается в девять утра и закрывается в шесть вечера.
the first expert prefers EnglishPythonmathsRussian

The colours tell the same story token by token. The Russian sentence goes to experts that prefer Russian almost everywhere, already in the first layer, where a token is still mostly what it looks like.

The Python line starts out mixed, many of its words sent to experts that prefer English, and ends almost entirely with Python’s experts, except its line breaks and indentation, which go, oddly, to experts that prefer Russian.

The maths problem ends with its own experts, except for the arithmetic, -, 7 and =, which goes to the experts that prefer code.

And the English sentence, in the last layer, sends more than half of its tokens to the experts of the maths problems, which are English too: by then the router sorts tokens less by language than by what should come next, and a sentence about opening hours reads a lot like a word problem.

5. A smaller model for one kind of text

If a domain uses only part of a layer’s experts, the rest are dead weight for it.

Pruning them is the cheapest way to shrink a mixture for one job: keep, in every layer, the \(N\) experts the domain used most on the calibration text, forbid the router to choose the others (their scores become \(-\infty\) before the softmax), and measure the loss on held-out text of the same kind and of the other kinds.

No training: whatever the pruned model can still do, it could do already.

def loss_on(batches):
    total, count = 0.0, 0
    with torch.no_grad():
        for i in range(0, len(batches), BATCH):
            ids = batches[i : i + BATCH].to(GPUS[0])
            logits = model(ids, causal_mask(ids))
            target = ids[:, 1:].to(logits.device)
            total += F.cross_entropy(
                logits[:, :-1].flatten(0, 1), target.flatten(), reduction="sum"
            ).item()
            count += target.numel()
    return total / count


def prune_for(domain, keep):
    d = DOMAINS.index(domain)
    for layer, block in enumerate(model.layers):
        mask = torch.zeros(n_experts, dtype=torch.bool)
        mask[np.argsort(-usage[d, layer])[:keep]] = True
        block.mlp.keep = mask.to(block.mlp.gate.weight.device)


def unprune():
    for block in model.layers:
        block.mlp.keep = None


KEEP = [n_experts, n_experts // 4 * 3, n_experts // 2, n_experts // 4]
TARGETS = ["Python", "Russian"]
results = []
start = time.time()
unprune()
full = {name: loss_on(data[name]["held out"]) for name in DOMAINS}
for target in TARGETS:
    for keep in KEEP[1:]:
        prune_for(target, keep)
        for name in DOMAINS:
            results.append(
                {
                    "pruned for": target,
                    "experts kept": keep,
                    "text": name,
                    "loss": loss_on(data[name]["held out"]),
                }
            )
unprune()
for target in TARGETS:
    for name in DOMAINS:
        results.append(
            {
                "pruned for": target,
                "experts kept": n_experts,
                "text": name,
                "loss": full[name],
            }
        )
results = pd.DataFrame(results)
print(f"{len(results)} evaluations in {time.time() - start:.0f} s")
table(
    results.pivot_table(
        index=["pruned for", "text"], columns="experts kept", values="loss"
    ).reset_index(),
    digits=3,
)

32 evaluations in 45 s
pruned fortext16324864
PythonEnglish8.8424.6723.2982.589
PythonPython3.8531.8951.5711.526
PythonRussian7.7493.5042.3691.861
Pythonmaths7.4213.1322.1451.960
RussianEnglish7.3193.7712.9092.589
RussianPython7.2512.6671.7691.526
RussianRussian4.4012.2031.9211.861
Russianmaths6.7612.8152.1511.960

The same numbers as a chart, with the memory of each pruned model in the second axis’ labels: the experts we keep, plus everything that isn’t an expert.

def gigabytes(keep):
    rest = total - parts["experts"]
    return (rest + parts["experts"] * keep / n_experts) * 2 / 1e9  # 16-bit weights


fig, axes = plt.subplots(1, len(TARGETS), figsize=(10, 3.6), sharey=True)
for ax, target in zip(axes, TARGETS):
    chunk = results[results["pruned for"] == target]
    for d, name in enumerate(DOMAINS):
        line = chunk[chunk.text == name].sort_values("experts kept")
        ax.plot(
            line["experts kept"],
            line["loss"],
            marker="o",
            color=palette[d],
            label=name,
            linewidth=2.4 if name == target else 1.2,
        )
    ax.set_xticks(KEEP, [f"{k}\n{gigabytes(k):.1f} GB" for k in KEEP], fontsize=8)
    ax.set_xlabel("experts kept per layer, and the model's size")
    ax.set_title(f"pruned for {target}", fontsize=10)
axes[0].set_ylabel("loss on held-out text")
axes[0].legend(frameon=False, fontsize=8.5)
plt.show()

Output of the code above

Pruning works in one direction only. Keep three quarters of the experts that Python uses, and Python barely notices, while every other kind of text gets visibly worse; the same holds for Russian.

At half the experts even the target domain pays a real price, and at a quarter the model falls apart for everything. So a mixture can be trimmed for a narrow job by a quarter without any training, which saves a quarter of the experts’ memory; cutting deeper would need fine-tuning to recover what is lost.

6. Experts in and out of GPU memory

A mixture saves compute, not memory: all 64 experts of every layer have to be somewhere, even though a token uses eight. Two numbers decide whether some of them can live elsewhere, in the CPU’s memory or on disk, and be fetched when needed.

The first is how many experts a batch of tokens touches. If routing were uniform, \(B\) tokens choosing \(k\) of \(E\) experts each would touch on average

\[E \left(1 - \left(1 - \frac{k}{E}\right)^{B}\right)\]

experts per layer, so a few dozen tokens in a batch already need nearly all of them. Real routing is not uniform, and we have the traces to measure it.

mixed = torch.cat(
    [traces[name] for name in DOMAINS]
)  # every routed token, of all four kinds
rng = np.random.default_rng(0)
sizes = [1, 2, 4, 8, 16, 32, 64, 128]
touched_real, touched_one = [], []
for b in sizes:
    picks = [mixed[rng.choice(len(mixed), b, replace=False)].long() for _ in range(50)]
    touched_real.append(
        np.mean(
            [
                [len(set(p[:, layer].flatten().tolist())) for layer in range(n_layers)]
                for p in picks
            ]
        )
    )
    one = traces["Python"]
    picks = [one[rng.choice(len(one), b, replace=False)].long() for _ in range(50)]
    touched_one.append(
        np.mean(
            [
                [len(set(p[:, layer].flatten().tolist())) for layer in range(n_layers)]
                for p in picks
            ]
        )
    )
k = olmoe_config["num_experts_per_tok"]
uniform = [n_experts * (1 - (1 - k / n_experts) ** b) for b in sizes]
fig, ax = plt.subplots(figsize=(7.2, 3.6))
ax.plot(
    sizes,
    uniform,
    color=COLORS["line"],
    linestyle="--",
    label="uniform routing, the formula",
)
ax.plot(
    sizes,
    touched_real,
    marker="o",
    color=COLORS["enc"],
    label="OLMoE, tokens of all four kinds",
)
ax.plot(
    sizes, touched_one, marker="o", color=COLORS["att"], label="OLMoE, Python tokens only"
)
ax.set_xscale("log", base=2)
ax.set_xticks(sizes, [str(s) for s in sizes])
ax.set_xlabel("tokens in the batch")
ax.set_ylabel("experts touched per layer")
ax.set_title(f"How many of the {n_experts} experts a batch needs")
ax.legend(frameon=False, fontsize=8.5)
plt.show()

Output of the code above

A batch of real tokens touches fewer experts than the formula predicts, because tokens of a kind prefer the same experts; a batch of Python alone touches fewer still.

But the curves join the formula’s soon enough: a few dozen tokens of mixed text already need most of every layer. That is why serving a mixture needs all its experts in fast memory as soon as more than a few users share a GPU.

DeepSeek goes the other way: its paper describes generation on 320 GPUs, “each GPU hosts only one expert”, so that every GPU gets a full batch of tokens for the one expert it holds[4].

The second number matters when a single user generates text one token at a time, as on a laptop: how often does the next token need an expert that the previous tokens already loaded?

We keep a cache of \(C\) experts per layer in GPU memory, evict the one used longest ago when a new one is needed, and count the hits along the real sequence of tokens.

def lru_hit_rate(trace, capacity):
    hits = total = 0
    for layer in range(trace.shape[1]):
        cache = OrderedDict()
        for chosen_experts in trace[:, layer].tolist():
            for e in chosen_experts:
                total += 1
                if e in cache:
                    hits += 1
                    cache.move_to_end(e)
                else:
                    cache[e] = True
                    if len(cache) > capacity:
                        cache.popitem(last=False)
    return hits / total


capacities = [8, 16, 24, 32, 48]
sequence = {
    name: traces[name][: 4 * LENGTH] for name in DOMAINS
}  # four windows in reading order
random_trace = torch.stack(
    [
        torch.stack([torch.randperm(n_experts)[:k] for _ in range(n_layers)])
        for _ in range(4 * LENGTH)
    ]
)
fig, ax = plt.subplots(figsize=(7.2, 3.6))
for d, name in enumerate(DOMAINS):
    ax.plot(
        capacities,
        [lru_hit_rate(sequence[name], c) for c in capacities],
        marker="o",
        color=palette[d],
        label=name,
    )
ax.plot(
    capacities,
    [lru_hit_rate(random_trace, c) for c in capacities],
    linestyle="--",
    color=COLORS["line"],
    label="uniform routing",
)
ax.set_xlabel(f"experts kept in GPU memory per layer, of {n_experts}")
ax.set_ylabel("share of expert uses already in memory")
ax.set_title("A least-recently-used cache of experts, one token at a time")
ax.set_ylim(0, 1)
ax.legend(frameon=False, fontsize=8.5)
plt.show()

A least-recently-used cache of experts, one token at a time

For a single user the picture is much better. Consecutive tokens of one text keep asking for the same experts, so a cache that holds a quarter of each layer’s experts already finds half to two thirds of what the next token needs, several times as often as random routing would; Russian and Python, with their small footprints, cache best.

This is what offloading schemes for running large mixtures on small GPUs rely on: the experts in use stay on the GPU, the rest wait in the CPU’s memory, and a miss costs one transfer of an expert’s weights.

7. Where mixtures are the right tool

A mixture of experts buys a big model’s knowledge at a small model’s compute per token, and pays in memory.

That trade suits a provider serving many users at once: the batch is large, every expert is busy, and the memory of all the experts is spread over enough GPUs that it doesn’t matter. This is why the largest open models of the last two years, DeepSeek-V3, Qwen3-235B, Kimi K2 and gpt-oss-120b, are mixtures.

On one GPU with one user the trade is worse. You hold 7 billion parameters to compute with 1.3, while a dense model of the same memory would put all of it to work for every token.

A mixture makes sense here when the memory is there anyway, when you can offload the experts a text doesn’t use, as section 6 showed works, or when you need the specialisation itself: section 5’s pruning gives a smaller model for one kind of text without training.

Fine-tuning a mixture has its own traps. The router can be trained or frozen; training it on a narrow dataset can undo the balance it learned, so freezing it is the safer default.

Adapters such as lesson 6’s LoRA can go on the attention and on the experts alike. Long context, lesson 9, and the memory of the KV cache, lesson 8, apply to a mixture unchanged: its attention is a dense model’s.

OLMoE’s passport, and the family tree with this lesson’s branch:

OLMoE's passport. Published September 2024 by the Allen Institute for AI, fully open: weights, data, code and logs, Apache 2.0. Type: decoder-only mixture of experts, 6.9 billion parameters, 1.3 billion active per token. Training: 5.1 trillion tokens, mostly English. Attention: 16 heads of 128, multi-head, QK-norm over the whole query and key projections. Positions: rotary, base 10,000. Experts: 64 SwiGLU experts of width 1,024 in each of 16 layers, 8 per token, no shared expert. Router: softmax over all experts, top 8, not renormalised; trained with a load-balancing loss and a z-loss. Vocabulary: 50,304 tokens, untied input and output. Good at: studying routing, English text at the cost of a 1B model. Not for: languages it barely saw, long contexts of more than 4,096 tokens OLMoE's passport. Published September 2024 by the Allen Institute for AI, fully open: weights, data, code and logs, Apache 2.0. Type: decoder-only mixture of experts, 6.9 billion parameters, 1.3 billion active per token. Training: 5.1 trillion tokens, mostly English. Attention: 16 heads of 128, multi-head, QK-norm over the whole query and key projections. Positions: rotary, base 10,000. Experts: 64 SwiGLU experts of width 1,024 in each of 16 layers, 8 per token, no shared expert. Router: softmax over all experts, top 8, not renormalised; trained with a load-balancing loss and a z-loss. Vocabulary: 50,304 tokens, untied input and output. Good at: studying routing, English text at the cost of a 1B model. Not for: languages it barely saw, long contexts of more than 4,096 tokens

The family of the Transformer as this course covers it, with lesson 7, mixtures of experts, highlighted and lessons 1 to 6 done The family of the Transformer as this course covers it, with lesson 7, mixtures of experts, highlighted and lessons 1 to 6 done

References

  1. Dai, D., Deng, C., Zhao, C., et al. (2024). DeepSeekMoE: towards ultimate expert specialization in mixture-of-experts language models. ACL 2024. ↩
  2. Muennighoff, N., Soldaini, L., Groeneveld, D., et al. (2024). OLMoE: open mixture-of-experts language models. ICLR 2025. ↩
  3. Cobbe, K., Kosaraju, V., Bavarian, M., et al. (2021). Training verifiers to solve math word problems. arXiv:2110.14168. Data: GSM8K, MIT licence. ↩
  4. DeepSeek-AI (2024). DeepSeek-V3 technical report. arXiv:2412.19437. ↩
  5. Shazeer, N., Mirhoseini, A., Maziarz, K., et al. (2017). Outrageously large neural networks: the sparsely-gated mixture-of-experts layer. ICLR 2017.
  6. Lepikhin, D., Lee, H., Xu, Y., et al. (2020). GShard: scaling giant models with conditional computation and automatic sharding. ICLR 2021.
  7. Fedus, W., Zoph, B., & Shazeer, N. (2021). Switch Transformers: scaling to trillion parameter models with simple and efficient sparsity. JMLR 2022.
  8. Zoph, B., Bello, I., Kumar, S., et al. (2022). ST-MoE: designing stable and transferable sparse expert models. arXiv:2202.08906.
  9. Jiang, A. Q., Sablayrolles, A., Roux, A., et al. (2024). Mixtral of experts. arXiv:2401.04088.
  10. Qwen Team (2025). Qwen3 technical report. arXiv:2505.09388.
  11. OpenAI (2025). gpt-oss-120b & gpt-oss-20b model card. arXiv:2508.10925.
]]> Eva Koroleva Modern LLM architecture and Qwen3 from scratch 2026-08-06T15:00:00+00:00 2026-11-08T00:00:00+00:00 https://blog.korolevalab.com/blog/2026/transformer-architectures-06-modern Transformer architectures: from BERT to diffusion · lesson 6 of 15 · all lessons · run on Kaggle · open in Colab · code on GitHub

Give Llama 3.1 8B a long contract to read, 128,000 tokens of it, and before it writes the first word of its answer it has to keep the whole contract in GPU memory: a key and a value for every token, in every head of every one of its 32 layers.

That memory, the KV cache, comes out larger than the model’s own weights. Its predecessor of the same width, Llama 2 7B, kept a key and a value for each of its 32 heads, four times as much per token.

Llama 3 keeps one pair for every four query heads, a change called grouped-query attention, and almost nobody trains a model without it any more.

Lesson 5 of the first course followed the changes Llama made to the original Transformer: RMSNorm, rotary positions, SwiGLU. This lesson follows what came after Llama, and there is a pattern in it.

Open the configuration file of any open model released since 2023 and nearly every setting that differs from Llama answers one of two problems: memory while the model generates, and stability while it trains.

The first gave us grouped-query attention, sliding windows and models that mix local and global layers. The second gave us QK-norm and soft-capping, small changes that keep attention from tearing a training run apart.

The decoder block of a modern LLM, with each change since Llama marked where it acts. Grouped-query attention shares one key-value head among several query heads and shrinks the KV cache. A sliding window limits how far back a local layer can look. QK-norm normalises queries and keys before their dot product, and soft-capping squashes the scores with a tanh. In most models the rest of the block is Llama's: RMSNorm before each part, rotary positions, a SwiGLU feed-forward network The decoder block of a modern LLM, with each change since Llama marked where it acts. Grouped-query attention shares one key-value head among several query heads and shrinks the KV cache. A sliding window limits how far back a local layer can look. QK-norm normalises queries and keys before their dot product, and soft-capping squashes the scores with a tanh. In most models the rest of the block is Llama's: RMSNorm before each part, rotary positions, a SwiGLU feed-forward network

What you’ll do

  1. Read the configuration files of a dozen models from Mistral 7B to gpt-oss and find what changed.
  2. Write grouped-query attention, find out what happens when a trained GPT-2 is forced to share its keys and values, and count the KV cache of every model in the table.
  3. Make small Transformers diverge on purpose, then keep them alive with QK-norm and soft-capping.
  4. Write Qwen3 in PyTorch, load the official weights of Qwen3-0.6B and match Hugging Face’s outputs, then let it think out loud and not.
  5. Give Qwen3 a sliding window it was never trained with, the way Mistral and Gemma build theirs, and see what it keeps and what it loses.
  6. Write LoRA from scratch and teach the 0.6B model to turn voice commands into JSON, in English and in Russian, against a model seven times its size.

What you need: a GPU. On Kaggle, Settings → Accelerator → GPU; on Colab, Runtime → Change runtime type → T4 GPU. The notebook downloads Qwen3-0.6B (1.5 GB), Qwen3-4B (8 GB), GPT-2 (550 MB) and the MASSIVE dataset (40 MB), trains a few small models and two LoRA adapters; on a T4 it runs in about 70 minutes.

Setup

One small helper package, tfs, draws the course’s plots and token chips and knows where the notebook runs (Kaggle, Colab or your own machine). Everything the lesson teaches is written out in the notebook itself.

import importlib.util
import subprocess
import sys

TFS = "git+https://github.com/xmarva/transformers-from-scratch@main"
if importlib.util.find_spec("tfs") is None:
    subprocess.check_call([sys.executable, "-m", "pip", "install", "-q", TFS])

from tfs import env
from tfs.viz import table, token_chips, use_course_style

env.quiet()  # no progress bars or environment warnings in the outputs
SEED = 42
env.seed_everything(SEED)
COLORS = use_course_style()
WORK_DIR = env.work_dir()

import json
import math
import random
import re
import tarfile
import time
import urllib.request

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import torch
import torch.nn as nn
import torch.nn.functional as F
from huggingface_hub import HfApi, hf_hub_download
from safetensors.torch import load_file
from transformers import AutoTokenizer
# the reference implementations, for the checks
from transformers import AutoModelForCausalLM, GPT2LMHeadModel

SMOKE = env.smoke()  # CI runs the lesson on a small slice; you run it in full
device = "cuda" if torch.cuda.is_available() else "cpu"
HALF = torch.float16 if device == "cuda" else torch.float32  # 16-bit matmuls crawl on a CPU
assert device == "cuda" or SMOKE, "This lesson trains models and runs a 4B one: use a GPU"
if device == "cuda":
    print("GPU:", torch.cuda.get_device_name(0))

GPU: Tesla T4

1. A dozen configs

Every model on Hugging Face comes with a config.json, a few dozen lines that say how wide the model is, how many layers and heads it has and which of the tricks of its year it uses.

It is the fastest way to read an architecture: the code behind it is usually shared with a dozen other models, and the config is what makes this one different. We download the configs of fourteen models released between September 2023 and August 2025.

The official repositories of Llama and Gemma ask you to accept a licence before you can download even the config, so for those we read the same files from Unsloth’s public copies.

MODELS = [  # repository, name, released; Llama and Gemma from Unsloth's public copies
    ("mistralai/Mistral-7B-v0.1", "Mistral 7B", "2023-09"),
    ("unsloth/gemma-2-9b", "Gemma 2 9B", "2024-06"),
    ("unsloth/Meta-Llama-3.1-8B", "Llama 3.1 8B", "2024-07"),
    ("unsloth/Llama-3.2-1B", "Llama 3.2 1B", "2024-09"),
    ("Qwen/Qwen2.5-7B", "Qwen2.5 7B", "2024-09"),
    ("allenai/OLMo-2-1124-7B", "OLMo 2 7B", "2024-11"),
    ("microsoft/phi-4", "Phi-4 14B", "2024-12"),
    ("mistralai/Mistral-Small-24B-Base-2501", "Mistral Small 24B", "2025-01"),
    ("unsloth/gemma-3-1b-pt", "Gemma 3 1B", "2025-03"),
    ("unsloth/gemma-3-27b-pt", "Gemma 3 27B", "2025-03"),
    ("Qwen/Qwen3-0.6B", "Qwen3 0.6B", "2025-04"),
    ("Qwen/Qwen3-8B", "Qwen3 8B", "2025-04"),
    ("HuggingFaceTB/SmolLM3-3B", "SmolLM3 3B", "2025-07"),
    ("openai/gpt-oss-20b", "gpt-oss 20B", "2025-08"),
]


def text_config(repo):
    with open(hf_hub_download(repo, "config.json")) as f:
        config = json.load(f)
    if "text_config" in config:  # a multimodal model keeps its language model here
        config = config["text_config"]
    if not config.get("head_dim"):
        config["head_dim"] = config["hidden_size"] // config["num_attention_heads"]
    return config


configs = {}
for repo, name, _ in MODELS:
    configs[name] = text_config(repo)
print(len(configs), "configs")

14 configs

Three settings need a word before the table. Some models give a layer either global attention, over everything before it, or local attention, over a window of the last few thousand tokens; section 7 is about why, and the config says it in one of three ways.

Soft-capping, one of the subjects of section 4, has a setting of its own, attn_logit_softcapping. QK-norm, the other, doesn’t appear in the config at all, because it is written into the model’s code: we note the model types whose code has it.

And Gemma’s configs don’t say whether the input and output embeddings are tied, because Gemma’s code ties them unless told otherwise.

# model types whose attention normalises queries and keys
QK_NORM = {"qwen3", "olmo2", "gemma3_text"}


def layer_kinds(c):
    n = c["num_hidden_layers"]
    if c.get("layer_types"):
        return ["local" if kind == "sliding_attention" else "global"
                for kind in c["layer_types"]]
    if c.get("sliding_window_pattern"):  # Gemma 3: every sixth layer is global
        pattern = c["sliding_window_pattern"]
        return ["global" if (i + 1) % pattern == 0 else "local" for i in range(n)]
    if c["model_type"] == "gemma2":  # Gemma 2 alternates, starting with a local layer
        return ["local" if i % 2 == 0 else "global" for i in range(n)]
    if c.get("sliding_window") and c.get("use_sliding_window", True):
        return ["local"] * n
    return ["global"] * n


def attention_summary(c):
    kinds = layer_kinds(c)
    local = kinds.count("local")
    if local == 0:
        return "all global"
    if local == len(kinds):
        return f"all local, {c['sliding_window']}"
    return f"{local} local : {len(kinds) - local} global, {c['sliding_window']}"


def tied(c):  # Gemma's configs leave the key out, and Gemma's code ties by default
    return c.get("tie_word_embeddings", c["model_type"].startswith("gemma"))


rows = []
for _, name, released in MODELS:
    c = configs[name]
    rows.append({
        "model": name,
        "released": released,
        "layers": c["num_hidden_layers"],
        "width": c["hidden_size"],
        "query heads": c["num_attention_heads"],
        "KV heads": c["num_key_value_heads"],
        "head size": c["head_dim"],
        "attention": attention_summary(c),
        "QK-norm": "yes" if c["model_type"] in QK_NORM else "",
        "soft-cap": c.get("attn_logit_softcapping") or "",
        "vocabulary": c["vocab_size"],
        "tied": "yes" if tied(c) else "",
    })
table(pd.DataFrame(rows))
modelreleasedlayerswidthquery headsKV headshead sizeattentionQK-normsoft-capvocabularytied
Mistral 7B2023-09324,096328128all local, 409632,000
Gemma 2 9B2024-06423,58416825621 local : 21 global, 409650.00256,000yes
Llama 3.1 8B2024-07324,096328128all global128,256
Llama 3.2 1B2024-09162,04832864all global128,256yes
Qwen2.5 7B2024-09283,584284128all global152,064
OLMo 2 7B2024-11324,0963232128all globalyes100,352
Phi-4 14B2024-12405,1204010128all global100,352
Mistral Small 24B2025-01405,120328128all global131,072
Gemma 3 1B2025-03261,1524125622 local : 4 global, 512yes262,144yes
Gemma 3 27B2025-03625,376321612852 local : 10 global, 1024yes262,208yes
Qwen3 0.6B2025-04281,024168128all globalyes151,936yes
Qwen3 8B2025-04364,096328128all globalyes151,936
SmolLM3 3B2025-07362,048164128all global128,256yes
gpt-oss 20B2025-08242,8806486412 local : 12 global, 128201,088

Read it column by column and the history of two years falls out.

KV heads is the first column that tells the models apart. Only OLMo 2 still has as many key-value heads as query heads; every other model shares one key-value head among two to eight query heads, and Gemma 3 1B goes all the way to a single key-value head for all four of its query heads. That is grouped-query attention, section 3.

Head size used to be the width divided by the number of heads. In Qwen3 0.6B it is not: sixteen heads of 128 make the queries 2,048 wide in a model only 1,024 wide. The width of attention became a setting of its own, chosen for the attention kernels and the KV cache, not derived from the rest.

Attention shows the second way to save memory, local layers. Mistral 7B made every layer local with a window of 4,096 tokens; its version 0.3 and Mistral Small went back to global attention everywhere, and Ministral 8B made three layers in four local again.

Gemma 2 alternates, Gemma 3 keeps one global layer in six, gpt-oss alternates with a window of only 128.

QK-norm and soft-cap are the stability tricks of section 4. Gemma 2 squashed its scores; Gemma 3, OLMo 2 and Qwen3 normalise queries and keys instead.

Vocabulary grew from Mistral’s 32,000 to over 262,000 in Gemma 3: more languages, fewer tokens per sentence, and in a small model a large share of all parameters. That is why small models often tie their input and output embeddings, using one matrix for both; Gemma ties them at every size.

A few settings didn’t fit in the table. SmolLM3 drops rotary positions from every fourth layer, which see the tokens with no positions at all; Gemma 3 27B scales its scores by the width per head of the whole model, 168, rather than by its head size, 128; the RoPE base, the number that sets the slowest rotation of rotary positions, went from 10,000 in Mistral 7B to 100 million in Mistral Small. Long context, where that number matters, is lesson 9.

print("SmolLM3 layers with rotary positions (1) and without (0):")
print(configs["SmolLM3 3B"]["no_rope_layers"])
gemma = configs["Gemma 3 27B"]
print("Gemma 3 27B scales scores by 1/sqrt of", gemma["query_pre_attn_scalar"],
      "with heads of", gemma["head_dim"])
for name in ["Mistral 7B", "Mistral Small 24B", "Llama 3.1 8B", "Qwen3 8B"]:
    print(f"{name}: RoPE base {configs[name]['rope_theta']:,.0f}")

SmolLM3 layers with rotary positions (1) and without (0):
[1, 1, 1, 0, 1, 1, 1, 0, 1, 1, 1, 0, 1, 1, 1, 0, 1, 1, 1, 0, 1, 1, 1, 0, 1, 1, 1, 0, 1, 1, 1, 0, 1, 1, 1, 0]
Gemma 3 27B scales scores by 1/sqrt of 168 with heads of 128
Mistral 7B: RoPE base 10,000
Mistral Small 24B: RoPE base 100,000,000
Llama 3.1 8B: RoPE base 500,000
Qwen3 8B: RoPE base 1,000,000

2. The decoder block they share

Under the settings, the fourteen models are variations of one machine, the one lesson 5 of the first course built step by step: a stack of decoder blocks with a residual connection around attention and around the feed-forward network, RMSNorm, rotary positions applied to queries and keys, a gated feed-forward network, and in most of them no biases at all.

Most put the norm before attention and before the feed-forward network and gate with SiLU, as Llama does. OLMo 2 normalises the outputs of the two instead of their inputs, Gemma 2 and 3 normalise both and gate with GELU, and gpt-oss replaces the feed-forward network with a mixture of experts, the subject of lesson 7.

GPT-2, in lesson 4, had LayerNorm, learned positions and GELU; of its block only the residual stream, the causal mask and, in most models, the norm before each part survived.

Two decoder blocks side by side. GPT-2's, from lesson 4: LayerNorm, attention with learned position embeddings added to the input, LayerNorm, a GELU feed-forward network, biases everywhere. The modern block: RMSNorm, attention with rotary positions applied to queries and keys inside it, RMSNorm, a SwiGLU feed-forward network with three matrices, no biases. The places where this lesson's changes act are marked: the key-value heads, the mask, the queries and keys, the scores Two decoder blocks side by side. GPT-2's, from lesson 4: LayerNorm, attention with learned position embeddings added to the input, LayerNorm, a GELU feed-forward network, biases everywhere. The modern block: RMSNorm, attention with rotary positions applied to queries and keys inside it, RMSNorm, a SwiGLU feed-forward network with three matrices, no biases. The places where this lesson's changes act are marked: the key-value heads, the mask, the queries and keys, the scores

What is left to change are the details of attention, and that is where the two pressures act. The rest of the lesson takes them in the order they arrive when you run a model: first generation, where memory runs out, then training, where stability does.

3. Grouped-query attention

Generating text is a strange workload for a GPU. To write one token the model multiplies a single vector by every weight matrix it has, and attends to every token before it.

The multiplications are trivial; moving the numbers is not. Each new token reads all the weights from memory once, which is shared by everyone in the batch, and reads its own sequence’s KV cache, which is shared by nobody.

In 2019 Noam Shazeer put it plainly: incremental inference “is often slow, due to the memory-bandwidth cost of repeatedly loading the large “keys” and “values” tensors” [1].

The cache for one token holds a key and a value of \(d_h\) numbers for each of \(n_{kv}\) heads in each of \(L\) layers, in \(b\) bytes per number:

\[\text{KV bytes per token} = 2 \cdot L \cdot n_{kv} \cdot d_h \cdot b.\]

In the original Transformer every query head had its own keys and values, so \(n_{kv} = n_h\). Shazeer’s multi-query attention kept all the query heads and gave them a single key-value head between them, \(n_{kv} = 1\): “the keys and values are shared across all of the different attention “heads””. The cache shrank by the number of heads, and decoding sped up, at the price of what Shazeer called “only minor quality degradation”. The authors of the next idea put it less kindly: MQA “can lead to quality degradation and training instability”.

In 2023 a Google team proposed the middle ground, grouped-query attention: split the \(n_h\) query heads into \(n_{kv}\) groups and give each group one key-value head[2]. Their verdict: GQA “achieves quality close to multi-head attention with comparable speed to MQA”.

Three ways to give query heads their keys and values, with eight query heads. Multi-head attention: each query head has its own key-value head, eight in all. Grouped-query attention: four groups of two query heads, each group sharing one key-value head, four in all. Multi-query attention: all eight query heads share one key-value head. Below each, the KV cache per token, in proportion: eight, four and one units Three ways to give query heads their keys and values, with eight query heads. Multi-head attention: each query head has its own key-value head, eight in all. Grouped-query attention: four groups of two query heads, each group sharing one key-value head, four in all. Multi-query attention: all eight query heads share one key-value head. Below each, the KV cache per token, in proportion: eight, four and one units

In code the sharing costs one line: before the scores are computed, each key-value head is repeated for the query heads of its group. Query head \(h\) reads key-value head \(\lfloor h / g \rfloor\), where \(g = n_h / n_{kv}\) is the size of a group.

def grouped_attention(q, k, v, mask):
    group = q.shape[1] // k.shape[1]  # query heads per key-value head
    k = k.repeat_interleave(group, dim=1)  # heads 0..g-1 read key-value head 0, and so on
    v = v.repeat_interleave(group, dim=1)
    scores = q @ k.transpose(-1, -2) / math.sqrt(q.shape[-1])
    scores = scores.masked_fill(~mask, float("-inf"))
    return scores.softmax(dim=-1) @ v


q = torch.randn(1, 8, 5, 16)  # batch, query heads, tokens, head size
k = torch.randn(1, 2, 5, 16)  # two key-value heads: groups of four
v = torch.randn(1, 2, 5, 16)
mask = torch.ones(5, 5, dtype=torch.bool).tril()
ours = grouped_attention(q, k, v, mask)
reference = F.scaled_dot_product_attention(q, k, v, attn_mask=mask, enable_gqa=True)
print("output", tuple(ours.shape))
print("max difference from PyTorch's GQA:", (ours - reference).abs().max().item())

output (1, 8, 5, 16)
max difference from PyTorch's GQA: 2.384185791015625e-07

The repeat makes a copy for clarity; fused attention kernels such as FlashAttention read the shared head in place, which is the point: in a group of \(g\), every key and value read from memory serves \(g\) query heads instead of one. The cache is \(g\) times smaller and each byte of it does \(g\) times more work.

Why not take an existing model and share its heads? Ainslie et al. did exactly that, but they then trained the converted model further with 5% of its original pre-training compute, and they had a reason.

Here is GPT-2 small, twelve heads in each of twelve layers, forced to share: within each group of heads we replace every head’s key and value projections by their average, the conversion the GQA paper found best, and measure the loss on Wikipedia text it has never seen. No training afterwards.

gpt2 = GPT2LMHeadModel.from_pretrained("openai-community/gpt2").to(device).eval()
gpt2_tok = AutoTokenizer.from_pretrained("openai-community/gpt2")
WIKITEXT = "Salesforce/wikitext"
path = hf_hub_download(WIKITEXT, "wikitext-103-raw-v1/test-00000-of-00001.parquet",
                       repo_type="dataset")
wiki = "".join(pd.read_parquet(path)["text"].tolist())
wiki_ids = gpt2_tok(wiki[:200_000]).input_ids[:8 * 1024]
wiki_batch = torch.tensor(wiki_ids, device=device).view(8, 1024)
original = {name: value.clone() for name, value in gpt2.state_dict().items()}


def share_heads(model, kv_heads):
    d = model.config.n_embd
    heads = model.config.n_head
    size = heads // kv_heads  # query heads per group
    for block in model.transformer.h:
        weight = block.attn.c_attn.weight.data  # columns: queries | keys | values
        bias = block.attn.c_attn.bias.data
        for part in [1, 2]:  # keys, then values
            columns = slice(part * d, (part + 1) * d)
            w = weight[:, columns].view(d, heads, d // heads)
            b = bias[columns].view(heads, d // heads)
            for g in range(kv_heads):
                group = slice(g * size, (g + 1) * size)
                w[:, group] = w[:, group].mean(dim=1, keepdim=True)
                b[group] = b[group].mean(dim=0, keepdim=True)


shared = []
for kv_heads in [12, 6, 4, 3, 2, 1]:
    gpt2.load_state_dict(original)
    share_heads(gpt2, kv_heads)
    with torch.no_grad():
        loss = gpt2(wiki_batch, labels=wiki_batch).loss.item()
    shared.append({"key-value heads": kv_heads, "query heads per group": 12 // kv_heads,
                   "loss": loss, "perplexity": math.exp(loss)})
table(pd.DataFrame(shared))
key-value headsquery heads per grouplossperplexity
1213.4932.65
626.59728.08
436.931,025.99
347.121,238.80
267.021,113.92
1127.201,342.50

Sharing ruins the model. With every head on its own, GPT-2 reads Wikipedia as well as it ever did.

Halve the key-value heads and its perplexity grows by more than an order of magnitude, and with fewer heads it stays about that bad. Each head of a trained model learned its own keys and values, and their average is none of them.

So a model is either trained with grouped queries from the start, as Llama 3 and most models after it were, or converted and then trained again.

Here is what it buys, for every model of section 1: the KV cache of one sequence of 32,768 tokens in 16-bit numbers, against the cache the same model would need with a key-value head for every query head. A local layer keeps only its window, and we count that too; section 7 explains it.

def kv_cache_bytes(c, tokens, kv_heads=None, windows=True):
    if kv_heads is None:
        kv_heads = c["num_key_value_heads"]
    total = 0
    for kind in layer_kinds(c):
        local = windows and kind == "local"
        kept = min(tokens, c["sliding_window"]) if local else tokens
        total += 2 * kept * kv_heads * c["head_dim"] * 2  # keys and values, 2 bytes each
    return total


context = 32_768
names = [name for _, name, _ in MODELS]
actual = [kv_cache_bytes(configs[n], context) / 1e9 for n in names]
full = []
for n in names:
    every_head = configs[n]["num_attention_heads"]
    full.append(kv_cache_bytes(configs[n], context, every_head, windows=False) / 1e9)
fig, ax = plt.subplots(figsize=(7.6, 5.2))
y = np.arange(len(names))
ax.barh(y, full, color=COLORS["grid"],
        label="a key-value head for every query head, all global")
ax.barh(y, actual, color=COLORS["att"], label="as configured")
for k in range(len(names)):
    ax.text(actual[k] + 0.4, k, f"{actual[k]:.1f} GB", va="center",
            color=COLORS["muted"], fontsize=9)
ax.set_yticks(y, names)
ax.invert_yaxis()
ax.set_xlabel(f"KV cache of one sequence of {context:,} tokens, GB")
ax.set_title("What grouped queries and local layers save")
ax.legend(frameon=False, loc="upper center", bbox_to_anchor=(0.5, -0.12), ncol=2)
ax.grid(axis="x")
ax.grid(axis="y", visible=False)
plt.show()

What grouped queries and local layers save

The grey bars are what these models would need with neither trick. Grouped queries alone cut the cache by the size of a group, two to eight times; local layers cut it further, and in Gemma 3 and gpt-oss the few global layers end up holding nearly all of it.

OLMo 2, with a key-value head for every query head and no local layers, needs the largest cache in the table, although it is far from the largest model.

And the claim this lesson opened with, checked: Llama 3.1 8B’s cache at its full context of 131,072 tokens against its weights, both in 16 bits.

llama = configs["Llama 3.1 8B"]
metadata = HfApi().get_safetensors_metadata("unsloth/Meta-Llama-3.1-8B")
llama_params = sum(metadata.parameter_count.values())
print(f"weights: {llama_params * 2 / 1e9:.1f} GB")
print(f"KV cache at 131,072 tokens: {kv_cache_bytes(llama, 131_072) / 1e9:.1f} GB")
print(f"the same with a key-value head per query head: "
      f"{kv_cache_bytes(llama, 131_072, llama['num_attention_heads']) / 1e9:.1f} GB")

weights: 16.1 GB
KV cache at 131,072 tokens: 17.2 GB
the same with a key-value head per query head: 68.7 GB

4. Training stability, QK-norm and soft-capping

The second pressure shows up during training, and it shows up late, which is what makes it expensive. In 2023 a Google team scaling a Vision Transformer to 22 billion parameters watched their 8-billion-parameter run diverge after a few thousand steps.

The cause was the attention scores: they grew without bound, the softmax over them turned into a one-hot vector, every head looked at a single token with an entropy near zero, and the loss went up [3].

The mechanism is easy to see in the formula. The score of query \(i\) for key \(j\) is

\[s_{ij} = \frac{q_i \cdot k_j}{\sqrt{d_h}}, \qquad q_i = W_Q x_i, \quad k_j = W_K x_j.\]

Both \(q\) and \(k\) grow with their weight matrices, so if training scales both \(W_Q\) and \(W_K\) by a factor \(c\), the scores grow by \(c^2\). A softmax doesn’t care about the scale of its inputs until it does: multiply scores of 1, 2 and 3 by 30 and almost all the weight lands on the largest. Past that point the gradient through the softmax vanishes for every other key, and a head that has collapsed onto one token has no way back.

Two fixes are in the table of section 1. QK-norm normalises each query and each key before the dot product, the same RMSNorm as everywhere else, with its own learned scale \(g\).

Qwen3 and Gemma 3 apply it to one head’s vector at a time, with one scale shared by the heads of a layer; OLMo 2 normalises the whole projection at once.

\[\hat q = g_q \odot \frac{q}{\text{rms}(q)}, \qquad \hat k = g_k \odot \frac{k}{\text{rms}(k)},\] \[s_{ij} = \frac{\hat q_i \cdot \hat k_j}{\sqrt{d_h}}.\]

Now \(\lVert \hat q \rVert = \sqrt{d_h}\) before the scale, and the score can only grow if training grows \(g_q\) and \(g_k\) themselves, slowly and in plain sight. Soft-capping, Gemma 2’s fix, leaves the scores alone until they get large and then bends them smoothly towards a ceiling \(c\):

\[s \;\leftarrow\; c \cdot \tanh\!\left(\frac{s}{c}\right).\]

Gemma 2 used \(c = 50\) in attention and \(c = 30\) on the final logits, the scores of the vocabulary[4].

s = np.linspace(-150, 150, 601)
fig, ax = plt.subplots(figsize=(6.4, 3.2))
ax.plot(s, s, color=COLORS["line"], linestyle="--", label="no cap")
ax.plot(s, 50 * np.tanh(s / 50), color=COLORS["att"],
        label="50 · tanh(s / 50), Gemma 2's attention")
ax.plot(s, 30 * np.tanh(s / 30), color=COLORS["dec"],
        label="30 · tanh(s / 30), its final logits")
ax.set_ylim(-80, 80)
ax.set_xlabel("score before the cap")
ax.set_ylabel("score after")
ax.set_title("Soft-capping bends large scores towards a ceiling")
ax.legend(frameon=False, loc="upper left")
plt.show()

Soft-capping bends large scores towards a ceiling

Gemma 3 dropped soft-capping in favour of QK-norm: “we replace the soft-capping of Gemma 2 with QK-norm”[5]. The two are not equivalent: a cap only limits the score, while QK-norm limits the vectors that make it, and leaves the model a learned scale to set how sharp attention may get.

You don’t need 8 billion parameters to watch this happen. Wortsman et al. showed that the same instabilities “also appear in small models when training at high learning rates”[6], and that turns a disaster costing millions into an experiment costing minutes.

Below is a small Transformer of the modern kind, six layers wide 384, reading Wikipedia byte by byte, with a switch for its attention: plain, with QK-norm, or with a soft cap at 50. It computes its scores by hand so that we can record the largest one in every layer as it trains.

class Norm(nn.Module):
    def __init__(self, dim):
        super().__init__()
        self.weight = nn.Parameter(torch.ones(dim))

    def forward(self, x):
        return self.weight * x * torch.rsqrt(x.pow(2).mean(dim=-1, keepdim=True) + 1e-6)


def rope_rotate(x, theta=10000.0):
    t, d = x.shape[-2], x.shape[-1]
    inv_freq = 1.0 / theta ** (torch.arange(0, d, 2, device=x.device).float() / d)
    angles = torch.arange(t, device=x.device).float()[:, None] * inv_freq
    angles = torch.cat([angles, angles], dim=-1)
    half = d // 2
    rotated = torch.cat([-x[..., half:], x[..., :half]], dim=-1)
    return x * angles.cos() + rotated * angles.sin()


class TinyAttention(nn.Module):
    def __init__(self, d, heads, variant):
        super().__init__()
        self.heads, self.head_dim, self.variant = heads, d // heads, variant
        self.qkv = nn.Linear(d, 3 * d, bias=False)
        self.out = nn.Linear(d, d, bias=False)
        if variant == "QK-norm":
            self.q_norm = Norm(self.head_dim)
            self.k_norm = Norm(self.head_dim)
        self.peak = 0.0

    def forward(self, x):
        b, t, d = x.shape
        q, k, v = self.qkv(x).view(b, t, 3, self.heads, self.head_dim).unbind(dim=2)
        if self.variant == "QK-norm":
            q, k = self.q_norm(q), self.k_norm(k)
        q = rope_rotate(q.transpose(1, 2))
        k = rope_rotate(k.transpose(1, 2))
        v = v.transpose(1, 2)
        scores = q @ k.transpose(-1, -2) / math.sqrt(self.head_dim)
        causal = torch.ones(t, t, dtype=torch.bool, device=x.device).tril()
        self.peak = scores.detach().masked_fill(~causal, 0).abs().amax()
        if self.variant == "soft-cap":
            scores = 50 * torch.tanh(scores / 50)
        weights = scores.masked_fill(~causal, float("-inf")).softmax(dim=-1)
        return self.out((weights @ v).transpose(1, 2).reshape(b, t, d))


class TinyBlock(nn.Module):
    def __init__(self, d, heads, variant):
        super().__init__()
        self.norm1, self.norm2 = Norm(d), Norm(d)
        self.attention = TinyAttention(d, heads, variant)
        hidden = 1024
        self.gate = nn.Linear(d, hidden, bias=False)
        self.up = nn.Linear(d, hidden, bias=False)
        self.down = nn.Linear(hidden, d, bias=False)

    def forward(self, x):
        x = x + self.attention(self.norm1(x))
        h = self.norm2(x)
        return x + self.down(F.silu(self.gate(h)) * self.up(h))


class TinyLM(nn.Module):
    def __init__(self, variant, d=384, layers=6, heads=6):
        super().__init__()
        self.embed = nn.Embedding(256, d)  # one row per byte
        self.blocks = nn.ModuleList([TinyBlock(d, heads, variant) for _ in range(layers)])
        self.norm = Norm(d)
        self.head = nn.Linear(d, 256, bias=False)
        for module in self.modules():
            if isinstance(module, (nn.Linear, nn.Embedding)):
                nn.init.normal_(module.weight, std=0.02)

    def forward(self, x):
        h = self.embed(x)
        for block in self.blocks:
            h = block(h)
        return self.head(self.norm(h))

The training loop is lesson 4’s: AdamW, a warm-up, a cosine decay, the gradient clipped at 1. The only thing we turn is the peak learning rate.

def wiki_bytes(split):
    name = f"wikitext-2-raw-v1/{split}-00000-of-00001.parquet"
    path = hf_hub_download(WIKITEXT, name, repo_type="dataset")
    text = "".join(pd.read_parquet(path)["text"].tolist())
    return torch.tensor(list(text.encode("utf-8")), dtype=torch.long)


train_bytes, valid_bytes = wiki_bytes("train"), wiki_bytes("validation")
STEPS = 30 if SMOKE else 600


def batch_of(data, generator, batch=32, length=256):
    starts = torch.randint(0, len(data) - length - 1, (batch,), generator=generator)
    x = torch.stack([data[s:s + length] for s in starts])
    y = torch.stack([data[s + 1:s + length + 1] for s in starts])
    return x.to(device), y.to(device)


def train_tiny(variant, peak_lr, steps=STEPS, warmup=100):
    torch.manual_seed(0)
    model = TinyLM(variant).to(device)
    optimizer = torch.optim.AdamW(model.parameters(), lr=peak_lr, betas=(0.9, 0.95),
                                  weight_decay=0.0)
    generator = torch.Generator().manual_seed(0)
    history = []
    for step in range(steps):
        if step < warmup:
            lr = peak_lr * (step + 1) / warmup
        else:
            progress = (step - warmup) / max(steps - warmup, 1)
            lr = peak_lr * (0.1 + 0.45 * (1 + math.cos(math.pi * progress)))
        for group in optimizer.param_groups:
            group["lr"] = lr
        x, y = batch_of(train_bytes, generator)
        loss = F.cross_entropy(model(x).flatten(0, 1), y.flatten())
        optimizer.zero_grad(set_to_none=True)
        loss.backward()
        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
        optimizer.step()
        if step % 5 == 0:
            peak = max(float(block.attention.peak) for block in model.blocks)
            history.append({"step": step, "loss": loss.item(), "peak score": peak})
    generator = torch.Generator().manual_seed(1)
    with torch.no_grad():
        losses = [F.cross_entropy(model(x).flatten(0, 1), y.flatten()).item()
                  for x, y in [batch_of(valid_bytes, generator) for _ in range(10)]]
    return pd.DataFrame(history), float(np.mean(losses))

RATES = [3e-3, 3e-2] if SMOKE else [1e-3, 1e-2, 1e-1]
SHOWN = RATES[1]  # the rate whose run the right-hand plot follows
VARIANTS = ["plain", "QK-norm", "soft-cap"]
runs = {}
start = time.time()
for variant in VARIANTS:
    for rate in RATES:
        runs[variant, rate] = train_tiny(variant, rate)
print(f"{len(runs)} runs in {time.time() - start:.0f} s")
results = [{"attention": v, "peak learning rate": r, "validation loss": runs[v, r][1]}
           for v in VARIANTS for r in RATES]
table(pd.DataFrame(results).pivot(index="peak learning rate", columns="attention",
                                  values="validation loss").reset_index(), digits=3)

9 runs in 1565 s
peak learning rateQK-normplainsoft-cap
0.0011.3801.4281.428
0.0101.4401.9941.570
0.1001.6222.6922.447
fig, axes = plt.subplots(1, 2, figsize=(10, 3.6))
for variant, color in zip(VARIANTS, [COLORS["dec"], COLORS["enc"], COLORS["att"]]):
    finals = [runs[variant, rate][1] for rate in RATES]
    axes[0].plot(RATES, finals, marker="o", color=color, label=variant)
    history = runs[variant, SHOWN][0]
    axes[1].plot(history["step"], history["peak score"], color=color, label=variant)
axes[0].set_xscale("log")
axes[0].set_xlabel("peak learning rate")
axes[0].set_ylabel("validation loss, nats per byte")
axes[0].set_title("Final loss against learning rate")
axes[1].set_yscale("log")
axes[1].set_xlabel("step")
axes[1].set_ylabel("largest attention score")
axes[1].set_title(f"The largest score at a learning rate of {SHOWN:g}")
axes[0].legend(frameon=False)
plt.show()

Final loss against learning rate

At the smallest learning rate the three models end close together. Raise it and they part: the plain model’s loss climbs fastest, soft-capping holds it back at the middle rate but not at the highest, and the model with QK-norm barely notices, and is a little better even at the lowest rate.

The right-hand plot shows why. In the plain model the largest score grows by orders of magnitude within the first couple of hundred steps and stays there, so its softmax is nearly one-hot. With QK-norm it stays in the tens.

The soft-capped model’s raw scores also grow far past the cap of 50, though less than the plain model’s, and the cap only squashes them on the way into the softmax: past the ceiling the tanh is flat and passes back almost no gradient.

5. Qwen3 in PyTorch

Qwen3 is a family of eight models that Alibaba’s Qwen team released in April 2025, from 0.6 to 235 billion parameters, six of them dense and two mixtures of experts, all under Apache 2.0 and trained on “36 trillion tokens” in “119 languages and dialects”[7].

The report describes the dense models in two sentences: like Qwen2.5 they use grouped-query attention, SwiGLU, rotary positions and RMSNorm with pre-normalisation, and then “we remove QKV-bias used in Qwen2 […] and introduce QK-Norm […] to the attention mechanism to ensure stable training for Qwen3”. That is the whole difference from the skeleton of section 2, and we already have every piece.

The smallest, Qwen3-0.6B, fits on a phone and runs comfortably on a T4. Its config:

REPO = "Qwen/Qwen3-0.6B"
config = text_config(REPO)
keys = ["hidden_size", "num_hidden_layers", "num_attention_heads", "num_key_value_heads",
        "head_dim", "intermediate_size", "vocab_size", "rope_theta", "rms_norm_eps",
        "tie_word_embeddings"]
for key in keys:
    print(f"{key:22} {config[key]}")

hidden_size            1024
num_hidden_layers      28
num_attention_heads    16
num_key_value_heads    8
head_dim               128
intermediate_size      3072
vocab_size             151936
rope_theta             1000000
rms_norm_eps           1e-06
tie_word_embeddings    True

The norm is RMSNorm, computed in 32-bit floats whatever the dtype of the model, as Qwen’s code does; the same class serves the residual stream and, applied to vectors of 128, the queries and keys.

class RMSNorm(nn.Module):
    def __init__(self, dim, eps):
        super().__init__()
        self.eps = eps
        self.weight = nn.Parameter(torch.ones(dim))

    def forward(self, x):
        dtype = x.dtype
        x = x.float()
        x = x * torch.rsqrt(x.pow(2).mean(dim=-1, keepdim=True) + self.eps)
        return self.weight * x.to(dtype)

Rotary positions as in lesson 5 of the first course, in the convention Qwen’s weights were trained with: the vector is split into halves rather than interleaved pairs, so dimension \(i\) is rotated together with dimension \(i + d_h/2\). With a base of a million, the slowest pair turns once every few million tokens. The tables of cosines and sines are computed once per forward pass, for the positions of every row of the batch.

def rope_tables(positions, head_dim, theta):
    exponents = torch.arange(0, head_dim, 2, device=positions.device).float() / head_dim
    inv_freq = 1.0 / theta ** exponents
    angles = positions.float()[..., None] * inv_freq
    angles = torch.cat([angles, angles], dim=-1)
    return angles.cos()[:, None], angles.sin()[:, None]  # shaped to broadcast over heads


def rotate(x, cos, sin):
    half = x.shape[-1] // 2
    rotated = torch.cat([-x[..., half:], x[..., :half]], dim=-1)
    return x * cos.to(x.dtype) + rotated * sin.to(x.dtype)

Attention puts sections 3 and 4 together. Queries take 16 heads of 128, keys and values 8: groups of two. QK-norm comes after the projection and before the rotation, one norm for all query heads and one for all key heads.

The cache stores keys and values with 8 heads, before they are repeated, which is where grouped queries save memory; the repeat happens on the way into the attention kernel. The attention itself is PyTorch’s scaled_dot_product_attention, which on a GPU usually runs a fused kernel that never writes out the matrix of scores.

class Cache:
    def __init__(self):
        self.k = None
        self.v = None

    def add(self, k, v):
        if self.k is not None:
            k = torch.cat([self.k, k], dim=2)
            v = torch.cat([self.v, v], dim=2)
        self.k, self.v = k, v
        return k, v


class Attention(nn.Module):
    def __init__(self, config):
        super().__init__()
        d = config["hidden_size"]
        self.heads = config["num_attention_heads"]
        self.kv_heads = config["num_key_value_heads"]
        self.head_dim = config["head_dim"]
        self.q_proj = nn.Linear(d, self.heads * self.head_dim, bias=False)
        self.k_proj = nn.Linear(d, self.kv_heads * self.head_dim, bias=False)
        self.v_proj = nn.Linear(d, self.kv_heads * self.head_dim, bias=False)
        self.o_proj = nn.Linear(self.heads * self.head_dim, d, bias=False)
        self.q_norm = RMSNorm(self.head_dim, config["rms_norm_eps"])
        self.k_norm = RMSNorm(self.head_dim, config["rms_norm_eps"])

    def forward(self, x, cos, sin, mask, cache=None):
        b, t, _ = x.shape
        q = self.q_proj(x).view(b, t, self.heads, self.head_dim)
        k = self.k_proj(x).view(b, t, self.kv_heads, self.head_dim)
        v = self.v_proj(x).view(b, t, self.kv_heads, self.head_dim)
        q = rotate(self.q_norm(q).transpose(1, 2), cos, sin)
        k = rotate(self.k_norm(k).transpose(1, 2), cos, sin)
        v = v.transpose(1, 2)
        if cache is not None:
            k, v = cache.add(k, v)
        group = self.heads // self.kv_heads
        k = k.repeat_interleave(group, dim=1)
        v = v.repeat_interleave(group, dim=1)
        out = F.scaled_dot_product_attention(q, k, v, attn_mask=mask)
        return self.o_proj(out.transpose(1, 2).reshape(b, t, -1))

The feed-forward network is SwiGLU with no biases, and the block is pre-norm with two residual connections. In Qwen3-0.6B, 1.7B and 4B the last step multiplies by the embedding matrix itself, because they tie their input and output embeddings; from 8B up the output matrix is a separate one, and the class takes either.

class FeedForward(nn.Module):
    def __init__(self, config):
        super().__init__()
        d, hidden = config["hidden_size"], config["intermediate_size"]
        self.gate_proj = nn.Linear(d, hidden, bias=False)
        self.up_proj = nn.Linear(d, hidden, bias=False)
        self.down_proj = nn.Linear(hidden, d, bias=False)

    def forward(self, x):
        return self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x))


class Block(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.input_layernorm = RMSNorm(config["hidden_size"], config["rms_norm_eps"])
        self.self_attn = Attention(config)
        d, eps = config["hidden_size"], config["rms_norm_eps"]
        self.post_attention_layernorm = RMSNorm(d, eps)
        self.mlp = FeedForward(config)

    def forward(self, x, cos, sin, mask, cache=None):
        x = x + self.self_attn(self.input_layernorm(x), cos, sin, mask, cache)
        return x + self.mlp(self.post_attention_layernorm(x))


class Qwen3(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.config = config
        self.embed_tokens = nn.Embedding(config["vocab_size"], config["hidden_size"])
        n_layers = config["num_hidden_layers"]
        self.layers = nn.ModuleList([Block(config) for _ in range(n_layers)])
        self.norm = RMSNorm(config["hidden_size"], config["rms_norm_eps"])
        self.tied = config.get("tie_word_embeddings", False)
        if not self.tied:  # Qwen3-8B and larger have an output matrix of their own
            d, vocab = config["hidden_size"], config["vocab_size"]
            self.lm_head = nn.Linear(d, vocab, bias=False)

    def hidden(self, ids, positions, mask, caches=None):
        x = self.embed_tokens(ids)
        c = self.config
        cos, sin = rope_tables(positions, c["head_dim"], c["rope_theta"])
        for i, layer in enumerate(self.layers):
            # one mask for all layers, or a list with one per layer (section 7)
            layer_mask = mask[i] if isinstance(mask, list) else mask
            cache = caches[i] if caches is not None else None
            x = layer(x, cos, sin, layer_mask, cache)
        return self.norm(x)

    def logits(self, h):
        if self.tied:
            return h @ self.embed_tokens.weight.T  # the output layer is the embeddings
        return self.lm_head(h)

    def forward(self, ids, positions, mask, caches=None):
        return self.logits(self.hidden(ids, positions, mask, caches))

6. The official weights in our code

We named our modules the way Qwen’s checkpoint names its tensors, so the transplant is a matter of dropping the prefix model.. The checkpoint of a tied model also stores lm_head.weight, an exact copy of the embedding matrix that the model never needs, and we skip it; load_state_dict refuses anything else missing or left over.

def load_qwen3(repo, dtype):
    config = text_config(repo)
    files = ["model.safetensors"]
    try:
        with open(hf_hub_download(repo, "model.safetensors.index.json")) as f:
            files = sorted(set(json.load(f)["weight_map"].values()))
    except Exception:
        pass  # a small model comes in one file
    state = {}
    for name in files:
        state.update(load_file(hf_hub_download(repo, name)))
    renamed = {}
    for key, value in state.items():
        if key == "lm_head.weight" and config.get("tie_word_embeddings", False):
            continue  # a copy of the embeddings
        renamed[key.removeprefix("model.")] = value.to(dtype)
    with torch.device("meta"):  # no memory for the random weights we are about to replace
        model = Qwen3(config)
    model.load_state_dict(renamed, strict=True, assign=True)
    return model.to(device).eval()


model = load_qwen3(REPO, torch.float32)
tok = AutoTokenizer.from_pretrained(REPO)
state = load_file(hf_hub_download(REPO, "model.safetensors"))
same = torch.equal(state["lm_head.weight"], state["model.embed_tokens.weight"])
print("lm_head equals the embeddings:", same)
del state

lm_head equals the embeddings: True

Where are the 596 million parameters? Not where you’d guess from a large model.

parts = {"embeddings (tied)": 0, "attention": 0, "QK-norm": 0, "feed-forward": 0,
         "other norms": 0}
for name, p in model.named_parameters():
    if name.startswith("embed_tokens"):
        parts["embeddings (tied)"] += p.numel()
    elif "q_norm" in name or "k_norm" in name:
        parts["QK-norm"] += p.numel()
    elif "self_attn" in name:
        parts["attention"] += p.numel()
    elif "mlp" in name:
        parts["feed-forward"] += p.numel()
    else:
        parts["other norms"] += p.numel()
total = sum(parts.values())
table(pd.DataFrame({"part": list(parts), "parameters": list(parts.values()),
                    "share": [f"{n / total:.1%}" for n in parts.values()]}))
partparametersshare
embeddings (tied)155,582,46426.1%
attention176,160,76829.6%
QK-norm7,1680.0%
feed-forward264,241,15244.3%
other norms58,3680.0%

About a quarter of the model is the embedding table, 151,936 rows of 1,024 numbers, one row per token, used twice: at the input and, tied, at the output.

Attention has a larger share than in older models because its heads are wider than the model itself. QK-norm adds a few thousand parameters, next to nothing for what section 4 showed it does.

The check, as in every lesson: the same tokens through our model and through Hugging Face’s, in 32-bit floats.

reference = AutoModelForCausalLM.from_pretrained(REPO, dtype=torch.float32)
reference = reference.to(device).eval()
text = ("The Rhine rises in the Swiss Alps, flows through Lake Constance and along "
        "the French–German border, and reaches the North Sea through a delta in the "
        "Netherlands.")
ids = tok(text, return_tensors="pt").input_ids.to(device)
t = ids.shape[1]
positions = torch.arange(t, device=device)[None]
causal = torch.ones(t, t, dtype=torch.bool, device=device).tril()
with torch.no_grad():
    ours = model(ids, positions, causal)
    theirs = reference(ids).logits
print("tokens:", t)
print("largest logit:", f"{theirs.abs().max().item():.2f}")
print("max difference:", f"{(ours - theirs).abs().max().item():.2e}")
del reference

tokens: 34
largest logit: 24.74
max difference: 6.96e-05

Our model and Hugging Face’s agree to within the rounding of 32-bit arithmetic, so the transplant is exact: the classes above are all there is to Qwen3.

Generation is lesson 4’s loop, extended to a batch of prompts of different lengths. The prompts are padded on the left, so that every row’s last token is in the same column; padding gets no attention, and positions are counted from each row’s first real token.

Every token is allowed to see itself, so that a row of pure padding never ends up with nothing to attend to, which would fill it with NaN.

@torch.no_grad()
def generate(model, prompts, max_new_tokens, stop_id, prefix=None):
    b = len(prompts)
    width = max(len(p) for p in prompts)
    ids = torch.full((b, width), stop_id, device=device)
    real = torch.zeros(b, width, dtype=torch.bool, device=device)
    for i, p in enumerate(prompts):
        ids[i, width - len(p):] = torch.tensor(p, device=device)
        real[i, width - len(p):] = True
    caches = [Cache() for _ in model.layers]
    start = 0
    if prefix is not None:  # section 8: keys and values of a shared start, computed once
        for cache, (k, v) in zip(caches, prefix["kv"]):
            cache.k, cache.v = k.expand(b, -1, -1, -1), v.expand(b, -1, -1, -1)
        start = prefix["length"]
        shared = torch.ones(b, start, dtype=torch.bool, device=device)
        real = torch.cat([shared, real], dim=1)
    positions = (real.long().cumsum(dim=1) - 1).clamp(min=0)[:, start:]
    causal = torch.ones(width, start + width, dtype=torch.bool, device=device)
    causal = causal.tril(diagonal=start)
    itself = torch.zeros(width, start + width, dtype=torch.bool, device=device)
    itself[torch.arange(width), start + torch.arange(width)] = True
    mask = (causal & real[:, None, None, :]) | itself
    h = model.hidden(ids, positions, mask, caches)
    next_ids = model.logits(h[:, -1]).argmax(dim=-1)
    written = [next_ids]
    done = next_ids == stop_id
    position = positions[:, -1]
    for _ in range(max_new_tokens - 1):
        if bool(done.all()):
            break
        position = position + 1
        real = torch.cat([real, torch.ones(b, 1, dtype=torch.bool, device=device)], dim=1)
        mask = real[:, None, None, :]
        h = model.hidden(next_ids[:, None], position[:, None], mask, caches)
        next_ids = model.logits(h[:, -1]).argmax(dim=-1)
        next_ids = torch.where(done, torch.full_like(next_ids, stop_id), next_ids)
        written.append(next_ids)
        done = done | (next_ids == stop_id)
    rows = []
    for row in torch.stack(written, dim=1).tolist():
        rows.append(row[:row.index(stop_id)] if stop_id in row else row)
    return rows

Qwen3 is one model with two modes. In thinking mode it first writes its reasoning between <think> and </think>, then the answer; in non-thinking mode the chat template inserts an empty <think></think> and the model answers at once. The template is part of the tokenizer, so the switch is one argument:

IM_END = tok.convert_tokens_to_ids("<|im_end|>")
question = ("Lena is three years older than Max. In five years their ages will add up "
            "to 41. How old is Max now?")
messages = [{"role": "user", "content": question}]
for thinking in [False, True]:
    prompt = tok.apply_chat_template(messages, tokenize=False, add_generation_prompt=True,
                                     enable_thinking=thinking)
    written = generate(model, [tok(prompt).input_ids], 2000 if thinking else 400, IM_END)[0]
    answer = tok.decode(written)
    print(f"--- thinking {'on' if thinking else 'off'}, {len(written)} tokens")
    print(answer if len(answer) < 1500 else answer[:700] + "\n[…]\n" + answer[-500:])

--- thinking off, 211 tokens
We are given:

- Lena is **three years older** than Max.
- In **five years**, their **ages will add up to 41**.

Let’s define:

- Let Max’s current age be $ x $.
- Then Lena’s current age is $ x + 3 $.

In **five years**, their ages will be:

- Max: $ x + 5 $
- Lena: $ x + 3 + 5 = x + 8 $

So, their ages in five years add up to:

$$
(x + 5) + (x + 8) = 2x + 13
$$

We are told this sum is 41:

$$
2x + 13 = 41
$$

Subtract 13 from both sides:

$$
2x = 28
$$

Divide by 2:

$$
x = 14
$$

### ✅ Final Answer:
Max is **14 years old now**.
--- thinking on, 1576 tokens
<think>
Okay, let's see. Lena is three years older than Max. So if I let Max's current age be something, then Lena's age would be Max's age plus three. Let me write that down.

Let me assign variables. Let M represent Max's current age. Then Lena's age would be M + 3. 

Now, the problem says that in five years, their ages will add up to 41. So I need to figure out their ages in five years and then set up an equation. 

In five years, Max's age will be M + 5, and Lena's age will be (M + 3) + 5. Let me calculate that. 

So adding those together: (M + 5) + (M + 3 + 5) = 41. Wait, maybe I should just add them directly. 

Alternatively, maybe I can write it as (M + 5) + (M + 3 + 5) = 41. Let me c
[…]
herefore, we can set up the equation:

$$
(M + 5) + (M + 8) = 41
$$

Simplifying:

$$
2M + 13 = 41
$$

Subtract 13 from both sides:

$$
2M = 28
$$

Divide by 2:

$$
M = 14
$$

---

**Verification:**

- Max's current age: $ 14 $
- Lena's current age: $ 14 + 3 = 17 $
- In five years: Max will be $ 14 + 5 = 19 $, and Lena will be $ 17 + 5 = 22 $
- Sum of their ages in five years: $ 19 + 22 = 41 $, which matches the problem statement.

---

**Final Answer:**

Max is currently $\boxed{14}$ years old.

Both answers are right. Without thinking the model writes a tidy solution at once. With thinking it first talks to itself for several times as many tokens, sets up the equation, doubts it, sets it up again and checks the result, and only then writes the same answer.

On a problem this easy the thinking buys nothing; it pays on problems where the first attempt goes wrong, and either way it is paid for in tokens.

7. Sliding-window attention and local layers

Grouped queries shrink the cache per token; a sliding window caps the number of tokens. In a local layer each token attends only to the last \(W\) tokens, itself included:

\[\text{token } i \text{ sees token } j \quad \Longleftrightarrow \quad i - W < j \le i.\]

Mistral 7B, in September 2023, gave every layer a window of 4,096 tokens and pointed out that the window is not a horizon: layer 2 sees layer 1’s outputs, each of which already summarises a window, so information travels about \(W\) tokens further back with every layer (\(W - 1\) with our mask; Mistral’s paper lets a position see the \(W\) tokens before itself). “At the last layer, using a window size of W=4096, we have a theoretical attention span of approximately 131K tokens”[8]. Its cache never holds more than \(W\) tokens: a rolling buffer writes token \(i\) into slot \(i \bmod W\), over the token that has just left the window, and on a 32k-token sequence “this reduces the cache memory usage by 8x, without impacting the model quality”.

Sliding-window attention. Left: the causal mask of global attention, a full lower triangle, against the mask of a local layer with a window of four, a band along the diagonal. Middle: through one, two and three local layers, what a token can draw on spans four, then seven, then ten tokens, itself included, because each layer reads the previous layer's summaries. Right: the rolling buffer, four slots, where token i is written into slot i mod 4 over the token that left the window Sliding-window attention. Left: the causal mask of global attention, a full lower triangle, against the mask of a local layer with a window of four, a band along the diagonal. Middle: through one, two and three local layers, what a token can draw on spans four, then seven, then ten tokens, itself included, because each layer reads the previous layer's summaries. Right: the rolling buffer, four slots, where token i is written into slot i mod 4 over the token that left the window

The rolling buffer is a few lines on top of our Cache. It keeps at most \(W\) keys and values; it doesn’t even need to keep them in order, because a token’s position is already rotated into its key and attention doesn’t care where in memory a key sits.

class RollingCache:
    def __init__(self, window):
        self.window = window
        self.k = None
        self.v = None
        self.seen = 0  # tokens written so far

    def add(self, k, v):
        for i in range(k.shape[2]):  # one token at a time, as during generation
            slot = self.seen % self.window
            if self.k is None:
                self.k = k[:, :, i:i + 1]
                self.v = v[:, :, i:i + 1]
            elif self.k.shape[2] < self.window:
                self.k = torch.cat([self.k, k[:, :, i:i + 1]], dim=2)
                self.v = torch.cat([self.v, v[:, :, i:i + 1]], dim=2)
            else:
                self.k[:, :, slot] = k[:, :, i]
                self.v[:, :, slot] = v[:, :, i]
            self.seen += 1
        return self.k, self.v

Two ways of computing the same thing must give the same answer: Qwen3 reading a text with a window of 16 in every layer, once with the full causal cache and a banded mask, once token by token with a rolling buffer of 16 slots and no mask at all. Qwen3 was not trained with a window, but the arithmetic doesn’t know that.

W = 16
t = ids.shape[1]
i = torch.arange(t, device=device)
band = (i[None, :] <= i[:, None]) & (i[:, None] - i[None, :] < W)
with torch.no_grad():
    banded = model(ids, positions, band)
    caches = [RollingCache(W) for _ in model.layers]
    stepwise = []
    for step in range(t):
        keys = min(step + 1, W)
        everything = torch.ones(1, 1, 1, keys, dtype=torch.bool, device=device)
        token, where = ids[:, step:step + 1], positions[:, step:step + 1]
        stepwise.append(model(token, where, everything, caches))
    stepwise = torch.cat(stepwise, dim=1)
print("tokens:", t, "window:", W, "slots in each cache:", caches[0].k.shape[2])
print("max difference:", f"{(banded - stepwise).abs().max().item():.2e}")

tokens: 34 window: 16 slots in each cache: 16
max difference: 1.01e-04

The two agree to rounding error, and the buffer never holds more than its 16 slots: a rolling cache is another way of storing what the banded mask lets each token see.

Gemma 2 made local attention common again: it alternated local and global layers, one to one.

Gemma 3 went to five local layers for every global one with a window of 1,024 (512 in the 1B model, says its config), Ministral 8B to three for every global one with a window of 32,768, and gpt-oss alternates again with a window of only 128 tokens.

The global layers carry information across the whole context; the local ones, half of the layers or more, cost a fixed amount of memory however long the context grows. Here are the layers of every model of section 1 that has local ones:

mixed = [name for name in names if "local" in layer_kinds(configs[name])]
longest = max(configs[n]["num_hidden_layers"] for n in mixed)
fig, ax = plt.subplots(figsize=(9.6, 0.42 * len(mixed) + 1.0))
for row, name in enumerate(mixed):
    for layer, kind in enumerate(layer_kinds(configs[name])):
        color = COLORS["att"] if kind == "global" else COLORS["att_tint"]
        cell = plt.Rectangle((layer, row - 0.36), 0.86, 0.72, color=color, linewidth=0)
        ax.add_patch(cell)
    window = configs[name]["sliding_window"]
    ax.text(longest + 0.8, row, f"window {window:,}", va="center",
            color=COLORS["muted"], fontsize=9)
ax.set_xlim(-0.5, longest + 9)
ax.set_ylim(len(mixed) - 0.4, -0.6)
ax.set_yticks(range(len(mixed)), mixed)
ax.set_xlabel("layer (dark: global attention, light: local)")
ax.set_title("Where the global layers are")
ax.grid(False)
ax.spines["left"].set_visible(False)
ax.spines["bottom"].set_visible(False)
plt.show()

Where the global layers are

How much does a model lose when most of its layers stop seeing far? Gemma’s authors answer for models trained that way: “We observe minimal impact on perplexity when changing this ratio”, and “the sliding window can be reduced significantly without impacting perplexity”.

We can ask a different question, of a model trained without any window: what happens to Qwen3 if we impose one?

We take a dozen Wikipedia articles of at least 4,096 tokens and measure the loss of predicting each token, averaged by its position, under several masks: global attention everywhere, as trained; a window of 512 in every layer; Gemma 3’s pattern, five local layers of 512 and then one global; and the same two with the first four tokens of the text kept visible to every layer, for a reason the results will make clear.

articles, current = [], []
for line in pd.read_parquet(path)["text"].tolist():
    # a new top-level heading, " = Title = ", starts an article
    if re.match(r"^ = [^=].* = \n$", line) and current:
        articles.append("".join(current))
        current = []
    current.append(line)
articles.append("".join(current))
LENGTH = 1024 if SMOKE else 4096
long_texts = []
for article in articles:
    article_ids = tok(article).input_ids
    if len(article_ids) >= LENGTH:
        long_texts.append(article_ids[:LENGTH])
long_texts = long_texts[:2 if SMOKE else 12]
print(len(long_texts), "articles of", LENGTH, "tokens")

12 articles of 4096 tokens
n_layers = config["num_hidden_layers"]
i = torch.arange(LENGTH, device=device)
causal = i[None, :] <= i[:, None]


def window_mask(window, sink=0):
    mask = causal & (i[:, None] - i[None, :] < window)
    if sink:
        mask = mask | (causal & (i[None, :] < sink))  # the first `sink` tokens stay visible
    return mask


def local_global(window, sink=0):
    local = window_mask(window, sink)
    return [causal if (k + 1) % 6 == 0 else local for k in range(n_layers)]


SETUPS = {
    "global everywhere (as trained)": [causal] * n_layers,
    "window of 512 everywhere": [window_mask(512)] * n_layers,
    "5 local : 1 global": local_global(512),
    "window of 512 + first 4 tokens": [window_mask(512, sink=4)] * n_layers,
    "5 local : 1 global + first 4 tokens": local_global(512, sink=4),
}
by_position = {}
for name, masks in SETUPS.items():
    total = torch.zeros(LENGTH - 1, device=device)
    for article_ids in long_texts:
        x = torch.tensor([article_ids], device=device)
        with torch.no_grad():
            h = model.hidden(x, i[None], masks)
            logits = model.logits(h[0, :-1]).float()
        total += F.cross_entropy(logits, x[0, 1:], reduction="none")
    by_position[name] = (total / len(long_texts)).cpu().numpy()

edges = [e for e in [0, 512, 1024, 2048] if e < LENGTH - 1] + [LENGTH - 1]
summary = []
for name, losses in by_position.items():
    row = {"attention": name}
    for lo, hi in zip(edges[:-1], edges[1:]):
        row[f"tokens {lo}–{hi}"] = losses[lo:hi].mean()
    summary.append(row)
table(pd.DataFrame(summary), digits=3)
attentiontokens 0–512tokens 512–1024tokens 1024–2048tokens 2048–4095
global everywhere (as trained)3.0662.9812.9352.812
window of 512 everywhere3.0666.3666.0856.009
5 local : 1 global3.0665.6045.7345.543
window of 512 + first 4 tokens3.0663.0583.1373.042
5 local : 1 global + first 4 tokens3.0663.0573.1273.030
fig, ax = plt.subplots(figsize=(8, 3.8))
colors = [COLORS["ink"], COLORS["dec"], COLORS["data"], COLORS["att"], COLORS["enc"]]
for (name, losses), color in zip(by_position.items(), colors):
    smooth = np.convolve(losses, np.ones(64) / 64, mode="valid")
    ax.plot(np.arange(len(smooth)) + 32, smooth, color=color, label=name, linewidth=1.4)
ax.axvline(512, color=COLORS["line"], linestyle="--")
ax.text(530, 0.5, "the window ends", color=COLORS["muted"], fontsize=9, va="center",
        transform=ax.get_xaxis_transform())
ax.set_xlabel("position of the predicted token")
ax.set_ylabel("loss, averaged over 64 positions")
ax.set_title("Qwen3-0.6B with windows it was never trained with")
ax.legend(frameon=False, fontsize=8.5, loc="upper center", bbox_to_anchor=(0.5, -0.2),
          ncol=2)
plt.show()

Qwen3-0.6B with windows it was never trained with

Up to the 512th token every mask is the same as global attention, and so is the loss. After it, a window of 512 in every layer is a disaster: the loss more than doubles and stays there.

One global layer in six helps only a little. The surprise is in the last two rows. Keep the first four tokens of the text visible to every layer, and the windowed model comes close to the global one all the way to the 4,096th token.

The first tokens matter not for what they say. Xiao et al. found that language models put a large share of their attention on the first tokens of a text, whatever those tokens are: a softmax has to give its weight to something, and the first tokens, visible to every later one, become the place to park it.

They called them attention sinks, and “keeping the KV of initial tokens will largely recover the performance of window attention”[9]. Cut the sinks off with a window and the heads that used them have nowhere to put their attention.

The gap that remains between the rows with sinks and global attention is what the model really loses by not seeing far, and it grows with position, as global attention gets more context to use.

Models trained with windows from the start learn to live with them, and gpt-oss builds the sink into every head: “Each attention head has a learned bias in the denominator of the softmax”[10].

8. Teaching the small model a job with LoRA

A model of 0.6 billion parameters will not write your novel, but it can do one narrow job fast and cheaply, on a laptop or a phone. The job here is the one every voice assistant does: turn a spoken command into a structured request. “Wake me up at five am this week” should become

{"intent": "alarm_set", "slots": [["time", "five am"], ["date", "this week"]]}

with the intent chosen from a fixed list and each slot’s value copied word for word from the command, so that the alarm app knows what to do.

The data is MASSIVE, which Amazon released in 2022: “1M realistic, parallel, labeled virtual assistant utterances” in 51 languages, with “18 domains, 60 intents, and 55 slots”[11].

Professional translators localised the English commands of an earlier dataset, SLURP, rather than translated them, so the Russian command about cultural events in California asks about the Leningrad region instead, and every slot is marked in the translation too. The licence is CC BY 4.0, and the official archive is 40 MB.

DATA = WORK_DIR / "massive"
if not (DATA / "1.1").exists():
    DATA.mkdir(parents=True, exist_ok=True)
    url = ("https://amazon-massive-nlu-dataset.s3.amazonaws.com/"
           "amazon-massive-dataset-1.1.tar.gz")
    urllib.request.urlretrieve(url, DATA / "massive.tar.gz")
    with tarfile.open(DATA / "massive.tar.gz") as archive:
        archive.extractall(DATA)


def read_locale(locale):
    rows = []
    for line in open(DATA / "1.1" / "data" / f"{locale}.jsonl", encoding="utf-8"):
        row = json.loads(line)
        # slots are marked inside the command: "wake me up at [time : five am]"
        row["slots"] = [[kind, value.strip()] for kind, value
                        in re.findall(r"\[(\w+) : ([^\]]+)\]", row["annot_utt"])]
        rows.append(row)
    return rows


english, russian = read_locale("en-US"), read_locale("ru-RU")
russian_by_id = {row["id"]: row for row in russian}
for part in ["train", "dev", "test"]:
    count = sum(row["partition"] == part for row in english)
    print(part, count, "commands in each language")
intents = sorted({row["intent"] for row in english if row["partition"] == "train"})
slot_types = sorted({kind for row in english if row["partition"] == "train"
                     for kind, _ in row["slots"]})
print(len(intents), "intents,", len(slot_types), "slot types")

train 11514 commands in each language
dev 2033 commands in each language
test 2974 commands in each language
60 intents, 55 slot types
def target(row):
    return json.dumps({"intent": row["intent"], "slots": row["slots"]}, ensure_ascii=False)


random.seed(SEED)
candidates = [row for row in english
              if row["partition"] == "test" and len(row["slots"]) >= 2]
shown = random.sample(candidates, 4)
table(pd.DataFrame({"English": [row["utt"] for row in shown],
                    "Russian": [russian_by_id[row["id"]]["utt"] for row in shown],
                    "English JSON": [target(row) for row in shown],
                    "Russian JSON": [target(russian_by_id[row["id"]]) for row in shown]}))
EnglishRussianEnglish JSONRussian JSON
convert current time from g. m. t. to e. s. t.конвертировть текущее время со среднего времени по гринвичу в североамериканское восточное время{"intent": "datetime_convert", "slots": [["time_zone", "g. m. t."], ["time_zone", "e. s. t."]]}{"intent": "datetime_convert", "slots": [["time_zone", "среднего времени по гринвичу"], ["time_zone", "североамериканское восточное время"]]}
play lady by kenny rogerвключи покинула чат клавы коки{"intent": "play_music", "slots": [["song_name", "lady"], ["artist_name", "kenny roger"]]}{"intent": "play_music", "slots": [["song_name", "покинула чат"], ["artist_name", "клавы коки"]]}
add a lunch with mikeдобавь обед с михаилом{"intent": "calendar_set", "slots": [["meal_type", "lunch"], ["person", "mike"]]}{"intent": "calendar_set", "slots": [["meal_type", "обед"], ["person", "михаилом"]]}
block one hour from ten tomorrow morningвыдели один час с десяти завтрашним утром{"intent": "calendar_set", "slots": [["time", "one hour"], ["time", "ten"], ["date", "tomorrow"], ["timeofday", "morning"]]}{"intent": "calendar_set", "slots": [["time", "один час"], ["time", "десяти"], ["date", "завтрашним"], ["timeofday", "утром"]]}

The localisation shows in every row. Kenny Rogers’s song became one by the Russian singer Klava Koka, and the slot keeps her name in the genitive the sentence puts it in, клавы коки; Mike became Mikhail, in the instrumental case, михаилом.

So a Russian slot value is not always in the form a dictionary gives, and a model has to copy it in the form it has in the command. The commands were written by people, typos included: конвертировть is missing a letter.

The score has three numbers, as in the MASSIVE paper: the share of commands whose intent is right, the F1 of the slots, counting a slot as right only if both its type and its value match, and the share of commands where everything matches exactly, which is what the alarm app needs. An answer that isn’t valid JSON gets nothing.

def parse(text):
    try:
        answer = json.loads(text.strip())
    except json.JSONDecodeError:
        return None
    if not isinstance(answer, dict) or not isinstance(answer.get("intent"), str):
        return None
    slots = []
    for slot in answer.get("slots", []) if isinstance(answer.get("slots"), list) else []:
        is_pair = isinstance(slot, list) and len(slot) == 2
        if is_pair and all(isinstance(s, str) for s in slot):
            slots.append((slot[0], slot[1].strip().lower()))
    return {"intent": answer["intent"], "slots": slots}


def score(rows, outputs):
    valid = intent = exact = 0
    found = wrong = missed = 0
    for row, output in zip(rows, outputs):
        gold = sorted((kind, value.lower()) for kind, value in row["slots"])
        answer = parse(output)
        if answer is None:
            missed += len(gold)
            continue
        valid += 1
        predicted = sorted(answer["slots"])
        intent += answer["intent"] == row["intent"]
        exact += answer["intent"] == row["intent"] and predicted == gold
        remaining = list(gold)
        for slot in predicted:
            if slot in remaining:
                found += 1
                remaining.remove(slot)
            else:
                wrong += 1
        missed += len(remaining)
    precision = found / max(found + wrong, 1)
    recall = found / max(found + missed, 1)
    f1 = 2 * precision * recall / max(precision + recall, 1e-9)
    n = len(rows)
    return {"valid JSON": valid / n, "intent": intent / n, "slot F1": f1,
            "exact": exact / n}


print(score(shown, [target(row) for row in shown]))
print(score(shown, ['{"intent": "alarm_set", "slots": []}'] * 4))

{'valid JSON': 1.0, 'intent': 1.0, 'slot F1': 1.0, 'exact': 1.0}
{'valid JSON': 1.0, 'intent': 0.0, 'slot F1': 0.0, 'exact': 0.0}

The baseline: ask nicely. Without any training, a chat model can be told the task in its system prompt: the format, the 60 intents, the 55 slot types and three examples. That prompt runs to a few hundred tokens, as the next cell prints, and it is the same for every command.

Recomputing its keys and values for every request would be a waste, so we compute them once and start every batch from a copy: prefix caching, which serving engines such as vLLM do automatically, in blocks of tokens, whenever prompts share a start. Our generate already takes a prefix; here is how to make one.

def schema_prompt(train_rows):
    examples = random.Random(0).sample([row for row in train_rows if row["slots"]], 3)
    lines = ["You turn commands for a voice assistant into JSON.",
             "Answer with one JSON object only: "
             '{"intent": INTENT, "slots": [[SLOT_TYPE, VALUE], ...]}.',
             "INTENT is exactly one of: " + ", ".join(intents) + ".",
             "SLOT_TYPE is exactly one of: " + ", ".join(slot_types) + ".",
             "VALUE is copied word for word from the command, in the command's language. "
             "Use an empty list when there are no slots.",
             "Examples:"]
    for row in examples:
        lines.append(row["utt"] + " -> " + target(row))
    return "\n".join(lines)


def chat_parts(system):
    marker = "<<COMMAND>>"
    messages = [{"role": "user", "content": marker}]
    if system:
        messages = [{"role": "system", "content": system}] + messages
    text = tok.apply_chat_template(messages, tokenize=False, add_generation_prompt=True,
                                   enable_thinking=False)
    return text.split(marker)  # what comes before the command, and after it


@torch.no_grad()
def make_prefix(model, prefix_ids):
    caches = [Cache() for _ in model.layers]
    t = len(prefix_ids)
    x = torch.tensor([prefix_ids], device=device)
    causal = torch.ones(t, t, dtype=torch.bool, device=device).tril()
    model.hidden(x, torch.arange(t, device=device)[None], causal, caches)
    return {"kv": [(cache.k, cache.v) for cache in caches], "length": t}


def run(model, rows, system=None, batch=64):
    before, after = chat_parts(system)
    prefix = make_prefix(model, tok(before, add_special_tokens=False).input_ids)
    prompts = [tok(row["utt"] + after, add_special_tokens=False).input_ids for row in rows]
    # similar lengths in a batch, so less padding
    order = sorted(range(len(rows)), key=lambda k: len(prompts[k]))
    outputs = [None] * len(rows)
    start = time.time()
    for s in range(0, len(order), batch):
        chosen = order[s:s + batch]
        answers = generate(model, [prompts[k] for k in chosen], 110, IM_END, prefix)
        for k, written in zip(chosen, answers):
            outputs[k] = tok.decode(written)
    return outputs, time.time() - start

The baselines run on a fixed sample of 500 test commands in each language, the same commands in English and in Russian. Qwen3-0.6B and Qwen3-4B, the latter in 16-bit floats so that it fits on a T4: our class loads it unchanged, since the dense models of the family share one architecture.

train_en = [row for row in english if row["partition"] == "train"]
train_ru = [row for row in russian if row["partition"] == "train"]
test_en = [row for row in english if row["partition"] == "test"]
test_ru = [russian_by_id[row["id"]] for row in test_en]
if SMOKE:
    test_en, test_ru = test_en[:32], test_ru[:32]
sample = sorted(random.Random(7).sample(range(len(test_en)), min(500, len(test_en))))
sample_en = [test_en[k] for k in sample]
sample_ru = [test_ru[k] for k in sample]
system = schema_prompt(train_en)
print(len(tok(system).input_ids), "tokens in the system prompt")

scores = []
outputs = {}
model = model.to(HALF)
baselines = [(REPO, 64)] if SMOKE else [(REPO, 64), ("Qwen/Qwen3-4B", 16)]
for repo, batch in baselines:
    if repo != REPO:
        del model
        torch.cuda.empty_cache()
        model = load_qwen3(repo, HALF)
    name = repo.split("/")[1]
    for language, rows in [("English", sample_en), ("Russian", sample_ru)]:
        written, seconds = run(model, rows, system, batch)
        method = name.removeprefix("Qwen3-") + ", instructions"
        outputs[method, language] = written
        scores.append({"method": method, "language": language, "commands": len(rows),
                       **score(rows, written), "seconds": seconds})
table(pd.DataFrame(scores), digits=3)

582 tokens in the system prompt
methodlanguagecommandsvalid JSONintentslot F1exactseconds
0.6B, instructionsEnglish5000.9900.4600.1700.04856.494
0.6B, instructionsRussian5000.9240.3460.0300.01073.703
4B, instructionsEnglish5000.9980.6940.3440.184148.817
4B, instructionsRussian5001.0000.6320.2670.128173.983

Instructions alone don’t make either model a voice assistant. Both write valid JSON nearly every time, so the format is not the problem; the content is. The 0.6B model picks the right intent for fewer than half of the commands and gets few slots right.

The 4B model does much better on intents, yet fewer than one command in five comes out exactly right, it takes several times as long, and both models do worse in Russian.

LoRA. Fine-tuning all 596 million parameters would work, but AdamW keeps two numbers per parameter on top of the parameter and its gradient, and a model per task is a lot of models. In 2021 a Microsoft team hypothesised that the change fine-tuning makes to a weight matrix has a low “intrinsic rank”, and proposed to learn only that change, as a product of two thin matrices[12]. A frozen weight \(W \in \mathbb{R}^{d_{out} \times d_{in}}\) becomes

\[W' = W + \frac{\alpha}{r} B A,\] \[A \in \mathbb{R}^{r \times d_{in}}, \qquad B \in \mathbb{R}^{d_{out} \times r},\]

with a rank \(r\) of 8 or 16 where \(d_{in}\) and \(d_{out}\) are in the thousands. \(B\) starts at zero, so the model starts exactly where it was; the factor \(\alpha / r\) saves retuning the learning rate when the rank changes.

For GPT-3 the paper reported that LoRA “can reduce the number of trainable parameters by 10,000 times and the GPU memory requirement by 3 times”, and since \(BA\) can be added into \(W\) after training, the fine-tuned model runs exactly as fast as the original.

LoRA. The input x goes through the frozen weight matrix W, d_out by d_in, and in parallel through two small trained matrices: A, which maps it down to r numbers, and B, which maps them back up to d_out. The two outputs are added, the low-rank path scaled by alpha over r. B starts at zero, so at the start the model is unchanged. After training, B times A is added into W and the extra path disappears LoRA. The input x goes through the frozen weight matrix W, d_out by d_in, and in parallel through two small trained matrices: A, which maps it down to r numbers, and B, which maps them back up to d_out. The two outputs are added, the low-rank path scaled by alpha over r. B starts at zero, so at the start the model is unchanged. After training, B times A is added into W and the extra path disappears

class LoRA(nn.Module):
    def __init__(self, base, rank=16, alpha=32):
        super().__init__()
        self.base = base  # the frozen linear layer
        where = base.weight.device
        self.A = nn.Parameter(torch.empty(rank, base.in_features, device=where))
        nn.init.kaiming_uniform_(self.A, a=math.sqrt(5))  # as PEFT; the paper: Gaussian
        self.B = nn.Parameter(torch.zeros(base.out_features, rank, device=where))
        self.scale = alpha / rank

    def forward(self, x):
        low_rank = (x @ self.A.T.to(x.dtype)) @ self.B.T.to(x.dtype)
        return self.base(x) + low_rank * self.scale


TARGETS = ["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"]


def add_lora(model):
    for p in model.parameters():
        p.requires_grad = False
    for block in model.layers:
        for part in [block.self_attn, block.mlp]:
            for name in TARGETS:
                if hasattr(part, name):
                    setattr(part, name, LoRA(getattr(part, name)))
    return [p for p in model.parameters() if p.requires_grad]

The training examples are chat turns with no system prompt at all: the command as the user’s message, the JSON as the assistant’s answer, followed by <|im_end|> so that the model learns to stop. The loss counts only the answer’s tokens.

The frozen model stays in 16-bit floats; the LoRA matrices, their gradients and AdamW’s state are 32-bit, and a gradient scaler keeps the 16-bit backward pass from underflowing. We train for one pass over the 11,514 English training commands, 32 to a step.

The 32 go through the model eight at a time and their gradients add up before the step: the step is the same, and long commands don’t run a T4 out of memory.

def examples_of(rows):
    before, after = chat_parts(None)
    examples = []
    for row in rows:
        prompt = tok(before + row["utt"] + after, add_special_tokens=False).input_ids
        answer = tok(target(row), add_special_tokens=False).input_ids + [IM_END]
        examples.append((prompt, answer))
    return examples


def answer_loss(model, chunk):
    """The loss summed over the answers' tokens of a few examples, padded to one width."""
    n = len(chunk)
    width = max(len(p) + len(a) for p, a in chunk)
    ids = torch.full((n, width), IM_END, device=device)
    labels = torch.full((n, width), -100, device=device)
    real = torch.zeros(n, width, dtype=torch.bool, device=device)
    for k, (prompt, answer) in enumerate(chunk):
        sequence = prompt + answer
        ids[k, :len(sequence)] = torch.tensor(sequence, device=device)
        # the logits at position p predict the token at p + 1
        first, last = len(prompt) - 1, len(sequence) - 1
        labels[k, first:last] = torch.tensor(answer, device=device)
        real[k, :len(sequence)] = True
    positions = torch.arange(width, device=device)[None].expand(n, -1)
    causal = torch.ones(width, width, dtype=torch.bool, device=device).tril()
    itself = torch.eye(width, dtype=torch.bool, device=device)
    mask = (causal & real[:, None, None, :]) | itself
    h = model.hidden(ids, positions, mask)
    scored = labels != -100
    logits = model.logits(h[scored]).float()
    return F.cross_entropy(logits, labels[scored], reduction="sum")


def train_lora(model, rows, batch=4 if SMOKE else 32, micro=2 if SMOKE else 8,
               peak_lr=2e-4, warmup=30):
    params = add_lora(model)
    optimizer = torch.optim.AdamW(params, lr=peak_lr, weight_decay=0.0)
    scaler = torch.amp.GradScaler(enabled=device == "cuda")
    examples = examples_of(rows)
    order = list(range(len(examples)))
    random.Random(1).shuffle(order)
    if SMOKE:
        order = order[:4 * batch]
    history = []
    model.train()
    for step, s in enumerate(range(0, len(order), batch)):
        chunk = [examples[k] for k in order[s:s + batch]]
        tokens = sum(len(answer) for _, answer in chunk)
        for group in optimizer.param_groups:
            group["lr"] = peak_lr * min(1.0, (step + 1) / warmup)
        optimizer.zero_grad(set_to_none=True)
        total = 0.0
        for m in range(0, len(chunk), micro):  # the pieces' gradients add up
            loss = answer_loss(model, chunk[m:m + micro]) / tokens
            scaler.scale(loss).backward()
            total += loss.item()
        scaler.unscale_(optimizer)
        torch.nn.utils.clip_grad_norm_(params, 1.0)
        scaler.step(optimizer)
        scaler.update()
        history.append(total)
    model.eval()
    return params, history

del model
torch.cuda.empty_cache()
model = load_qwen3(REPO, HALF)
start = time.time()
params, history = train_lora(model, train_en)
trained = sum(p.numel() for p in params)
frozen = sum(p.numel() for p in model.parameters()) - trained
print(f"{len(history)} steps in {time.time() - start:.0f} s")
print(f"trained {trained:,} parameters, {trained / frozen:.2%} of the frozen {frozen:,}")

fig, ax = plt.subplots(figsize=(6.4, 3.0))
ax.plot(history, color=COLORS["enc"], linewidth=1)
ax.set_yscale("log")
ax.set_xlabel("step (32 commands each)")
ax.set_ylabel("loss on the answer's tokens")
ax.set_title("LoRA on Qwen3-0.6B trained on English commands only")
plt.show()

360 steps in 328 s
trained 10,092,544 parameters, 1.69% of the frozen 596,049,920

LoRA on Qwen3-0.6B trained on English commands only

The adapter has never seen a Russian command, and it gets no instructions: just the command, as in training.

def evaluate(model, method):
    for language, rows in [("English", test_en), ("Russian", test_ru)]:
        written, seconds = run(model, rows)
        outputs[method, language] = written
        scores.append({"method": method, "language": language, "commands": len(rows),
                       **score(rows, written), "seconds": seconds})


evaluate(model, "0.6B + LoRA, en")
table(pd.DataFrame(scores), digits=3)
methodlanguagecommandsvalid JSONintentslot F1exactseconds
0.6B, instructionsEnglish5000.9900.4600.1700.04856.494
0.6B, instructionsRussian5000.9240.3460.0300.01073.703
4B, instructionsEnglish5000.9980.6940.3440.184148.817
4B, instructionsRussian5001.0000.6320.2670.128173.983
0.6B + LoRA, enEnglish2,9741.0000.8540.7680.655121.610
0.6B + LoRA, enRussian2,9740.9990.7640.4300.396155.080

One pass over the English commands, a few minutes on a T4 and ten million trained parameters change the picture. On the whole English test set the adapted 0.6B model gets about two commands in three exactly right, several times as many as the 4B model given instructions, and it answers faster.

The more surprising row is Russian. The adapter never saw a Russian command, yet the model gets a good share of them exactly right: the skill it learned, find the intent and copy the words, carries over to a language Qwen3 learned in pretraining.

For scale, the MASSIVE paper’s baselines, multilingual models of a few hundred million parameters fine-tuned in full, reach 69.7–72.5% exact match in English when trained on all 51 languages, and 42.6–52.8% in Russian when trained on English alone (its Table 7).

After one pass, with a small fraction of its parameters trained, our model lands a few points below the first range and somewhat below the second.

And the promise that LoRA costs nothing at inference, checked: add \(\frac{\alpha}{r}BA\) into every frozen matrix, drop the extra path, and the logits stay the same up to 16-bit rounding.

def merge(model):
    for block in model.layers:
        for part in [block.self_attn, block.mlp]:
            for name in TARGETS:
                layer = getattr(part, name, None)
                if isinstance(layer, LoRA):
                    with torch.no_grad():
                        delta = (layer.B.float() @ layer.A.float()) * layer.scale
                        layer.base.weight += delta.to(layer.base.weight.dtype)
                    setattr(part, name, layer.base)


probe = tok(chat_parts(None)[0] + "wake me up at five am this week" + chat_parts(None)[1],
            add_special_tokens=False).input_ids
x = torch.tensor([probe], device=device)
t = x.shape[1]
causal = torch.ones(t, t, dtype=torch.bool, device=device).tril()
with torch.no_grad():
    before_merge = model(x, torch.arange(t, device=device)[None], causal).float()
    merge(model)
    after_merge = model(x, torch.arange(t, device=device)[None], causal).float()
print("parameters now:", f"{sum(p.numel() for p in model.parameters()):,}")
print("largest logit:", f"{before_merge.abs().max().item():.1f}")
difference = (before_merge - after_merge).abs().max().item()
print("max difference after merging:", f"{difference:.3f}")

parameters now: 596,049,920
largest logit: 35.3
max difference after merging: 0.080

After the merge the model has exactly as many parameters as before training, and its logits move only by what rounding the merged weights to 16 bits does. The adapter is now part of the weights, and the model runs as fast as it did before training.

Now the same recipe with the Russian training commands added, 23,028 commands in all, still one pass.

del model
torch.cuda.empty_cache()
model = load_qwen3(REPO, HALF)
start = time.time()
params, history_both = train_lora(model, train_en + train_ru)
print(f"{len(history_both)} steps in {time.time() - start:.0f} s")
evaluate(model, "0.6B + LoRA, en+ru")
table(pd.DataFrame(scores), digits=3)

720 steps in 664 s
methodlanguagecommandsvalid JSONintentslot F1exactseconds
0.6B, instructionsEnglish5000.9900.4600.1700.04856.494
0.6B, instructionsRussian5000.9240.3460.0300.01073.703
4B, instructionsEnglish5000.9980.6940.3440.184148.817
4B, instructionsRussian5001.0000.6320.2670.128173.983
0.6B + LoRA, enEnglish2,9741.0000.8540.7680.655121.610
0.6B + LoRA, enRussian2,9740.9990.7640.4300.396155.080
0.6B + LoRA, en+ruEnglish2,9741.0000.8690.7740.677122.265
0.6B + LoRA, en+ruRussian2,9741.0000.8510.7140.622144.666
methods = ["0.6B, instructions", "4B, instructions",
           "0.6B + LoRA, en", "0.6B + LoRA, en+ru"]
bar_colors = [COLORS["line"], COLORS["line"], COLORS["enc"], COLORS["enc"]]
fig, axes = plt.subplots(1, 2, figsize=(10, 3.4), sharey=True)
for ax, language, rows in zip(axes, ["English", "Russian"], [sample_en, sample_ru]):
    values = []
    for method in methods:
        written = outputs.get((method, language))
        if written is not None and len(written) > len(rows):  # scored on the whole test set
            written = [written[k] for k in sample]
        values.append(score(rows, written)["exact"] if written is not None else 0)
    ax.barh(methods, values, color=bar_colors)
    for k, value in enumerate(values):
        ax.text(value + 0.01, k, f"{value:.0%}", va="center", color=COLORS["muted"],
                fontsize=9)
    ax.set_xlim(0, 1)
    ax.set_title(f"Exactly right on the same {len(rows)} {language} commands")
    ax.grid(axis="x")
    ax.grid(axis="y", visible=False)
axes[0].invert_yaxis()
plt.show()

Output of the code above

The Russian commands lift the Russian score most of the way to the English one, and English loses nothing; if anything it gains a little. That brings the 0.6B model close to the paper’s models trained on all 51 languages, which reach 66.2–69.2% exact match in Russian, after a few minutes of training on a free GPU.

The mistakes that are left are worth reading, especially in Russian, where a slot’s value is a word in some case of a declension and has to be copied in exactly that case.

mistakes = []
for row, written in zip(test_ru, outputs["0.6B + LoRA, en+ru", "Russian"]):
    if score([row], [written])["exact"] < 1:
        mistakes.append({"command": row["utt"], "correct": target(row), "model": written})
random.Random(3).shuffle(mistakes)
table(pd.DataFrame(mistakes[:8]))
commandcorrectmodel
у меня встреча в полдень сегодня{"intent": "calendar_set", "slots": [["event_name", "встреча"], ["timeofday", "полдень"], ["date", "сегодня"]]}{"intent": "calendar_query", "slots": [["event_name", "встреча"], ["timeofday", "полдень"], ["date", "сегодня"]]}
музыка король и шут{"intent": "play_music", "slots": [["artist_name", "король и шут"]]}{"intent": "play_music", "slots": [["song_name", "король и шут"]]}
закажи пиццу с колбасками в доминос пицце{"intent": "takeaway_order", "slots": [["food_type", "пиццу"], ["ingredient", "колбасками"], ["business_name", "доминос пицце"]]}{"intent": "takeaway_order", "slots": [["food_type", "пиццу с колбасками"], ["business_name", "доминос пицце"]]}
какие магазины рядом со мной{"intent": "recommendation_locations", "slots": [["business_type", "магазины"]]}{"intent": "recommendation_locations", "slots": []}
где алла пугачева покупает одежду своим детям{"intent": "general_quirky", "slots": [["person", "алла пугачева"]]}{"intent": "qa_factoid", "slots": [["person", "алла пугачева"]]}
сколько антон готов вложить в мою компанию на собрании которое состоится сегодня в пять двадцать вечера{"intent": "calendar_query", "slots": [["person", "антон"], ["event_name", "собрании"], ["date", "сегодня"], ["time", "пять двадцать вечера"]]}{"intent": "social_post", "slots": [["person", "антон"], ["event_name", "собрании"], ["time", "пять двадцать вечера"]]}
восход в москве по владивостокскому времени{"intent": "datetime_convert", "slots": [["time", "восход"], ["place_name", "москве"], ["place_name", "владивостокскому"]]}{"intent": "datetime_convert", "slots": [["place_name", "восход"], ["place_name", "москве"], ["place_name", "владивостокскому"]]}
хочу побольше узнать о ярмарке на этих выходных{"intent": "general_quirky", "slots": [["date", "на этих выходных"]]}{"intent": "recommendation_events", "slots": [["event_name", "ярмарке"], ["date", "на этих выходных"]]}

None of these mistakes is about grammar. The model copies every value in the form the command has it, колбасками, москве, владивостокскому. What it gets wrong is the annotation.

Some are plain errors: Король и Шут is a band, not a song, and магазины should have been marked as a type of business. Some are boundaries: is the order пиццу с колбасками, or a pizza plus an ingredient?

And some are judgement calls where the reference itself is debatable: does “у меня встреча в полдень сегодня” add a meeting or ask about one, and is a question about where Alla Pugacheva buys her children’s clothes a quirky remark or a factual question?

Part of the distance between two thirds right and all right is disagreement of this kind, which no model will remove.

9. Where these models are the right tool

Every model in this lesson is a decoder, and by now they differ less in what they are than in how they spend memory. That makes the config the first thing to read when you choose one.

Two models with the same number of parameters can need very different amounts of memory per conversation: count the key-value heads, the head size and the local layers, as section 3 did, and you know how many users and how much context fit on your GPU before you download a single weight.

Lesson 8 does that sum for a T4 and adds DeepSeek’s latent attention, which compresses the cache further still.

For a narrow job with a fixed format, such as extraction, routing, classification with a reason, or turning text into a call to your own API, a small model with LoRA is the cheapest thing that works.

Here it beat a model seven times its size that was given the task in words, and it learned the Russian version of the job from a few minutes of examples.

The adapter is a few dozen megabytes, so one base model in memory can serve many jobs by swapping adapters, and after merging it runs as fast as the base model.

For anything that needs knowledge or open reasoning, size matters more than any trick in this lesson. A 0.6B model knows what fits in 0.6 billion numbers, and thinking mode buys accuracy with tokens, which on a long problem means thousands of them.

The tricks themselves are now the default, and you will rarely choose them: nearly every new model shares its keys and values among query heads, by grouped queries or, in DeepSeek’s models, by compressing them; QK-norm is common, and local layers are spreading.

What this lesson left out is the next step in each direction. gpt-oss and the larger Qwen3 models are mixtures of experts, which spend only a fraction of their parameters on each token: lesson 7. Long context, and the RoPE bases of a hundred million in the table of section 1, are lesson 9.

Qwen3’s passport, and the family tree with this lesson’s branch:

Qwen3's passport. Published April 2025 by Alibaba's Qwen team, eight models under Apache 2.0. Type: decoder-only, six dense models and two mixtures of experts. Training: 36 trillion tokens in 119 languages, thinking and non-thinking modes in one model. Attention: grouped queries, 16 query heads and 8 key-value heads of 128 in the 0.6B model, with QK-norm. Positions: rotary with a base of one million, no sliding window. Normalisation: pre-norm RMSNorm, and RMSNorm on every query and key head. Activation: SwiGLU, 3,072 wide in the 0.6B model, no biases. Vocabulary: byte-level BPE with 151,936 rows; small models tie input and output embeddings. Sizes: 0.6B, 1.7B, 4B, 8B, 14B and 32B dense, 30B-A3B and 235B-A22B mixtures of experts. Good at: running locally, narrow jobs after LoRA, reasoning in thinking mode. Not for: facts it can't look up Qwen3's passport. Published April 2025 by Alibaba's Qwen team, eight models under Apache 2.0. Type: decoder-only, six dense models and two mixtures of experts. Training: 36 trillion tokens in 119 languages, thinking and non-thinking modes in one model. Attention: grouped queries, 16 query heads and 8 key-value heads of 128 in the 0.6B model, with QK-norm. Positions: rotary with a base of one million, no sliding window. Normalisation: pre-norm RMSNorm, and RMSNorm on every query and key head. Activation: SwiGLU, 3,072 wide in the 0.6B model, no biases. Vocabulary: byte-level BPE with 151,936 rows; small models tie input and output embeddings. Sizes: 0.6B, 1.7B, 4B, 8B, 14B and 32B dense, 30B-A3B and 235B-A22B mixtures of experts. Good at: running locally, narrow jobs after LoRA, reasoning in thinking mode. Not for: facts it can't look up

The family of the Transformer as this course covers it, with lesson 6, modern LLMs, highlighted and lessons 1 to 5 done The family of the Transformer as this course covers it, with lesson 6, modern LLMs, highlighted and lessons 1 to 5 done

References

  1. Shazeer, N. (2019). Fast Transformer decoding: one write-head is all you need. arXiv:1911.02150. ↩
  2. Ainslie, J., Lee-Thorp, J., de Jong, M., Zemlyanskiy, Y., Lebrón, F., & Sanghai, S. (2023). GQA: training generalized multi-query Transformer models from multi-head checkpoints. EMNLP 2023. ↩
  3. Dehghani, M., Djolonga, J., Mustafa, B., et al. (2023). Scaling vision Transformers to 22 billion parameters. ICML 2023. ↩
  4. Gemma Team (2024). Gemma 2: improving open language models at a practical size. arXiv:2408.00118. ↩
  5. Gemma Team (2025). Gemma 3 technical report. arXiv:2503.19786. ↩
  6. Wortsman, M., Liu, P. J., Xiao, L., et al. (2023). Small-scale proxies for large-scale Transformer training instabilities. ICLR 2024. ↩
  7. Qwen Team (2025). Qwen3 technical report. arXiv:2505.09388. ↩
  8. Jiang, A. Q., Sablayrolles, A., Mensch, A., et al. (2023). Mistral 7B. arXiv:2310.06825. ↩
  9. Xiao, G., Tian, Y., Chen, B., Han, S., & Lewis, M. (2023). Efficient streaming language models with attention sinks. ICLR 2024. ↩
  10. OpenAI (2025). gpt-oss-120b & gpt-oss-20b model card. arXiv:2508.10925. ↩
  11. FitzGerald, J., Hench, C., Peris, C., et al. (2022). MASSIVE: a 1M-example multilingual natural language understanding dataset with 51 typologically-diverse languages. ACL 2023. Data: github.com/alexa/massive, CC BY 4.0. ↩
  12. Hu, E. J., Shen, Y., Wallis, P., Allen-Zhu, Z., Li, Y., Wang, S., Wang, L., & Chen, W. (2021). LoRA: low-rank adaptation of large language models. ICLR 2022. ↩
  13. OLMo Team (2024). 2 OLMo 2 furious. arXiv:2501.00656.
]]> Eva Koroleva Theories of consciousness and the hard problem 2026-07-10T00:00:00+00:00 2026-11-08T00:00:00+00:00 https://blog.korolevalab.com/blog/2026/theories-of-consciousness

In short

“Consciousness is what the brain does” may well be true, and it still doesn’t explain why any of what the brain does is felt. Most philosophers lean physicalist, and most of them also think there is a hard problem.

The thought experiments, from Mary’s room to fading qualia, are tools for finding out which kind of physicalist you are. The scientific theories disagree about where in the brain consciousness lives, and the first preregistered head-to-head test, published in 2025, found key predictions of both leading theories wrong.

Meanwhile the question has become practical: in hospitals, where a quarter of patients who seem unresponsive turn out to follow instructions, and in AI.

62%of philosophers say there is a hard problem of consciousness; 52% lean physicalist
60 of 241patients with no response to commands performed mental tasks in a scanner
0 of 655prefrontal electrodes showed the ignition one leading theory predicted

In 1998 the neuroscientist Christof Koch bet the philosopher David Chalmers that within twenty-five years science would find the mechanism by which the brain’s neurons produce consciousness. On 23 June 2023, at a conference in New York, they agreed publicly that it hadn’t, and declared Chalmers the winner[1].

The occasion was the first results of an experiment designed to make that kind of progress.

The two most influential scientific theories of consciousness had agreed in advance on predictions where they disagreed, theory-neutral labs had recorded 256 people with fMRI, MEG and electrodes inside the skull, and both theories had come out with some of their own predictions failed[2].

If you work in machine learning or engineering, you probably think the question is settled in principle: consciousness is something the brain does, the details are for neuroscience, and philosophers who talk about a “hard problem” are making a mystery out of a gap in the data.

This post is written for you. It is not an argument against physicalism. It is an argument that physicalism is a family of positions that disagree about the most interesting question, and that the thought experiments philosophers use are the fastest way to find out which member of the family you are.

What is the hard problem of consciousness?

In 1974 Thomas Nagel gave the definition that most of the field still uses: “an organism has conscious mental states if and only if there is something that it is like to be that organism”[3]. His example was a bat.

Bats find their way by echolocation, and imagining hanging upside down in a cave and listening for echoes tells me what it would be like for me to behave like a bat, not what it is like for a bat to be a bat.

Every subjective fact is tied to a point of view, Nagel argued, while physical science works by abandoning points of view.

Nagel did not conclude that physicalism is false. He wrote: “It would be a mistake to conclude that physicalism must be false.” His point was that we don’t understand how it could be true, as a pre-Socratic philosopher who said that matter is energy would not have understood what he was saying.

In 1995 David Chalmers sharpened this into a distinction that gave the field its vocabulary[4]. The easy problems of consciousness are to explain functions: how we discriminate stimuli and react to them, integrate information, report our mental states, focus attention, control behaviour deliberately, and differ awake from asleep.

They are easy only in the sense that we know what an explanation would look like: a mechanism that performs the function. The hard problem is the question left after all of them are solved: “Why is the performance of these functions accompanied by experience?”

The easy problems of consciousness are to explain functions: discriminating and reacting to stimuli, integrating information, reporting mental states, focusing attention, deliberate control of behaviour, and the difference between waking and sleep. Each is explained by finding a mechanism that performs it. The hard problem is why the performance of these functions is accompanied by experience. Below, Levine's explanatory gap: heat equals molecular motion explains why heat melts ice, while pain equals C-fibre firing explains wincing but not why pain feels the way it does The easy problems of consciousness are to explain functions: discriminating and reacting to stimuli, integrating information, reporting mental states, focusing attention, deliberate control of behaviour, and the difference between waking and sleep. Each is explained by finding a mechanism that performs it. The hard problem is why the performance of these functions is accompanied by experience. Below, Levine's explanatory gap: heat equals molecular motion explains why heat melts ice, while pain equals C-fibre firing explains wincing but not why pain feels the way it does

Twelve years earlier Joseph Levine had put the same worry as a gap in explanation rather than in reality[5]. “Heat is the motion of molecules” explains: once you know how molecules move, you see why anything with that motion warms your hand and melts ice.

“Pain is the firing of C-fibres” explains why pain makes you wince and withdraw, and leaves something out: “what is left unexplained by the discovery of C-fiber firing is why pain should feel the way it does!” Levine was explicit that this doesn’t show materialism false. It shows that the identity, even if true, doesn’t explain.

One more distinction prevents a lot of confusion, especially in conversations about AI. Ned Block separated access consciousness, information that is available to reasoning and to the control of speech and action, from phenomenal consciousness: “Phenomenal consciousness is experience; the phenomenally conscious aspect of a state is what it is like to be in that state”[6].

Access is a function and could be engineered. Phenomenal consciousness is what the hard problem is about. Much of the talk about whether a model “is conscious” slides from one to the other without noticing.

Is consciousness just what the brain does?

The slogan can be true and still not do the job you want from it. Here is why, in four steps.

First, it states a thesis and doesn’t explain anything. When J. J. C. Smart defended the identity of sensations with brain processes in 1959, he offered it as a hypothesis supported by simplicity, not as an explanation of why brain processes feel like something[7].

Second, “what the brain does” is ambiguous between the two main kinds of physicalism. Is it the brain’s function, the pattern of causes and effects, which a silicon copy or a nation of people with radios could reproduce?

Or is it the brain’s biology, which they couldn’t? The slogan can’t decide whether a machine could be conscious, and that is the question everyone actually asks.

Third, physicalists disagree among themselves about what is left over once the functions are explained. Chalmers sorted them into types[8]. Type-A physicalists, like Daniel Dennett, say nothing is left: there is no hard problem beyond the easy ones.

Type-B physicalists say there is a real gap in our understanding but not in the world: experience is physical, and we just can’t see it a priori, as people couldn’t see a priori that water is H₂O. Type-C physicalists say the gap is real and science will close it.

Fourth, the philosophers who have thought hardest about it don’t find it obvious. In the 2020 PhilPapers survey of 1,785 professional philosophers, 52% accepted or leaned towards physicalism about the mind and 32% towards non-physicalism; asked whether there is a hard problem of consciousness, 62% said yes and 30% no[9]. Many physicalists, in other words, think the hard problem is real.

Four bar charts from the 2020 PhilPapers Survey. Mind: physicalism 52%, non-physicalism 32%. Is there a hard problem: yes 62%, no 30%. Zombies: conceivable but not metaphysically possible 37%, metaphysically possible 24%, inconceivable 16%. Consciousness: functionalism 33%, dualism 22%, identity theory 13%, panpsychism 8%, eliminativism 5% Four bar charts from the 2020 PhilPapers Survey. Mind: physicalism 52%, non-physicalism 32%. Is there a hard problem: yes 62%, no 30%. Zombies: conceivable but not metaphysically possible 37%, metaphysically possible 24%, inconceivable 16%. Consciousness: functionalism 33%, dualism 22%, identity theory 13%, panpsychism 8%, eliminativism 5%

The honest starting point is Nagel’s: physicalism may be true, and we don’t yet understand how. The rest of the debate is about what kind of explanation would close that gap, or whether it needs closing.

Thought experiments about consciousness

Thought experiments in philosophy of mind work like unit tests: each isolates one assumption and shows what you are committed to if you accept it. You don’t have to accept the conclusion to learn something from where you resist.

Mary’s room

Frank Jackson’s Mary is a brilliant scientist who has lived all her life in a black-and-white room, seeing the world through a black-and-white television[10]. She knows everything physical there is to know about colour vision: the wavelengths, the retina, the neurons, how it all leads someone to say “the sky is blue”.

In the 1986 version she knows all physical facts, completed physics and every functional role included[11]. Then she walks out and sees a red tomato. Does she learn something?

Mary's room. Mary has lived all her life in a black-and-white room and knows every physical fact about colour vision. She walks out and sees a red tomato. Does she learn something? Jackson in 1982 said yes, a new fact, so physicalism is false. The physicalist replies: she gains new abilities to imagine and recognise red (Lewis, Nemirow); she gains acquaintance (Conee); she learns an old fact under a new concept (Loar, Papineau). Jackson himself gave up the argument by 1995 Mary's room. Mary has lived all her life in a black-and-white room and knows every physical fact about colour vision. She walks out and sees a red tomato. Does she learn something? Jackson in 1982 said yes, a new fact, so physicalism is false. The physicalist replies: she gains new abilities to imagine and recognise red (Lewis, Nemirow); she gains acquaintance (Conee); she learns an old fact under a new concept (Loar, Papineau). Jackson himself gave up the argument by 1995

“It seems just obvious that she will learn something about the world and our visual experience of it,” Jackson wrote. If she does, her complete physical knowledge was not complete knowledge, and physicalism is false.

Physicalists have three main replies. She learns no new fact but gains new abilities, to imagine, remember and recognise red, as David Lewis and Laurence Nemirow argued. She becomes acquainted with something she knew about.

Or, the most common reply today, she learns an old physical fact under a new concept, a phenomenal concept that she can only have by having the experience, as you can know that the morning star is the evening star without being able to recognise it at dawn[12].

And then there is the twist that most popular accounts leave out: Jackson himself gave up the argument.

By 1995 he had concluded that if experiences were epiphenomenal, causing nothing, as his own version of the argument made them, nobody could know about them, so something had to be wrong with it; he later joined the physicalists[13].

Philosophical zombies

A philosophical zombie is not a robot that acts conscious. It is an exact physical duplicate of you, atom for atom, that behaves exactly as you do, says the same things, including “I’m conscious”, and has no experience at all: “There is nothing it is like to be a zombie”[14].

Nobody thinks zombies exist. The argument is about whether they are possible: if a zombie world is conceivable in a strict sense, then it is possible, and if it is possible, the physical facts don’t fix the facts about experience, and physicalism is false[15].

The zombie argument. You and your zombie twin have the same atoms and the same behaviour, and both say they are conscious; you have experience and the twin has none. Premise one: such a twin is conceivable. Premise two: what is conceivable is possible. Conclusion: the physical facts do not fix experience, so physicalism is false. In 2020, 16% of philosophers said zombies are inconceivable, 37% conceivable but not possible, 24% possible The zombie argument. You and your zombie twin have the same atoms and the same behaviour, and both say they are conscious; you have experience and the twin has none. Premise one: such a twin is conceivable. Premise two: what is conceivable is possible. Conclusion: the physical facts do not fix experience, so physicalism is false. In 2020, 16% of philosophers said zombies are inconceivable, 37% conceivable but not possible, 24% possible

The argument has an ancestor in Saul Kripke’s lectures of 1970. To make heat, God need only make molecules move; there is nothing more to heat than that, and the sensation of warmth is a separate thing.

To make pain, Kripke argued, God “must do something in addition to the mere creation of the C-fiber stimulation”: he must make it felt[16].

The survey shows where physicalists get off. In 2020, 16% of philosophers said zombies are inconceivable, the type-A answer; 37% said conceivable but not possible, the type-B answer; 24% said possible.

The modal physicalist grants that you can imagine a zombie and denies that imagining it tells you anything about what is possible, as you can imagine water that isn’t H₂O.

The inverted spectrum

John Locke asked whether a violet might produce in one person the idea that a marigold produces in another, with nobody ever able to tell, since both would call violets violet.

If such an inversion were possible between two people who behave identically, behaviourism would be false; between two people with identical functional organisation, functionalism; only between two physical duplicates would it threaten physicalism[17].

Here empirical science pushes back on a thought experiment. Colour space isn’t symmetric: there are more distinguishable shades between red and blue than between green and yellow, dark yellow is brown while dark blue is still blue, and red looks warm.

A naive red-green inversion would show up in behaviour. Defenders retreat to symmetric cases that are harder to picture.

The China brain

In 1978 Ned Block imagined converting the government of China to functionalism and having it run a human mind for an hour[18].

Each citizen gets a two-way radio and follows instructions that together implement the machine table describing your mind; letters displayed on satellites, visible across the country, keep track of which state the whole system is in.

If functionalism is right, the system has your mental states, including the felt ones. Is anybody there, besides a billion people with radios? (Block chose China because it “has a billion inhabitants”, not, as is often said, because that matches the number of neurons; a brain has about 86 billion.)

Most people’s intuition says no one is there. The functionalist replies that intuitions about huge, slow systems are unreliable, and that the mind would belong to the system, not to any citizen.

Fading and dancing qualia

The thought experiment that should interest engineers most comes from a dualist. Chalmers asks you to replace your neurons, one at a time, with silicon chips that do exactly the same local job (Chalmers, 1995b). Your behaviour never changes, because each chip does what the neuron did.

Suppose the end result, a silicon copy of your brain, had no experience. Then somewhere in the middle there is a person with half-faded experience, who sees pale pink where you saw bright red, while saying exactly what you would say about the bright red.

“The crucial point here is that Joe is systematically wrong about everything that he is experiencing.” Chalmers finds that implausible, and concludes that experience doesn’t fade: a fine-grained functional copy of a conscious brain is conscious, in silicon or anything else.

Fading qualia. Your neurons are replaced one at a time by silicon chips that do the same local job: none, a quarter, half, three quarters, all. Behaviour and speech stay exactly the same at every step. If the all-silicon brain had no experience, a middle person would see pale pink while sincerely describing bright red, systematically wrong about their own experience. Chalmers concludes experience does not fade: a fine-grained functional copy is conscious Fading qualia. Your neurons are replaced one at a time by silicon chips that do the same local job: none, a quarter, half, three quarters, all. Behaviour and speech stay exactly the same at every step. If the all-silicon brain had no experience, a middle person would see pale pink while sincerely describing bright red, systematically wrong about their own experience. Chalmers concludes experience does not fade: a fine-grained functional copy is conscious

The “dancing” version installs a silicon circuit that sees blue where you see red as a backup and flips a switch between it and your own neurons: “My experiences are switching from red to blue, but I do not notice any change.”

That is the standard philosophical argument that consciousness doesn’t depend on carbon, and its author believes experience is not physical. Being a dualist about what experience is and a functionalist about which systems have it is a consistent position.

What is it like to be a bat?

The bat belongs on this list because it tests something different: the limits of imagination rather than of physics. Learning everything about echolocation tells you what a bat does with sound. Whether that will ever tell you what it is like to be a bat is the question that keeps every other thought experiment alive.

Philosophical positions on consciousness

The thought experiments sort people into positions, and the positions are easiest to remember as answers to three questions.

A decision tree of positions on consciousness. First question: is there a hard problem beyond explaining the functions? No: type-A physicalism, illusionism and eliminativism (Dennett, Frankish). Yes: second question, is experience physical? Yes: type-B physicalism, an identity known only a posteriori (Loar, Papineau); biological naturalism (Searle); mysterianism, true but beyond us (McGinn). No: third question, then what is it? A separate thing that acts on the body: substance dualism (Descartes). Fundamental properties linked to the physical by laws: property dualism (Chalmers). Caused by the brain but causing nothing: epiphenomenalism. The inner nature of matter itself: Russellian monism and panpsychism (Russell, Strawson, Goff). Everything is mental: idealism (Kastrup). Functionalism, which systems are conscious, cuts across the tree A decision tree of positions on consciousness. First question: is there a hard problem beyond explaining the functions? No: type-A physicalism, illusionism and eliminativism (Dennett, Frankish). Yes: second question, is experience physical? Yes: type-B physicalism, an identity known only a posteriori (Loar, Papineau); biological naturalism (Searle); mysterianism, true but beyond us (McGinn). No: third question, then what is it? A separate thing that acts on the body: substance dualism (Descartes). Fundamental properties linked to the physical by laws: property dualism (Chalmers). Caused by the brain but causing nothing: epiphenomenalism. The inner nature of matter itself: Russellian monism and panpsychism (Russell, Strawson, Goff). Everything is mental: idealism (Kastrup). Functionalism, which systems are conscious, cuts across the tree

Illusionism and eliminativism

Illusionism, defended by Keith Frankish and in its classic form by Daniel Dennett, is the most radical physicalism and the most often misunderstood. “This is the view that phenomenal consciousness, as usually conceived, is illusory”[19].

Illusionists don’t deny that you see, feel pain or have experiences in the everyday sense. They deny that experiences have the special properties philosophers attribute to them, intrinsic, private, ineffable qualities, and say that introspection misrepresents ordinary physical states as having them.

The hard problem becomes the illusion problem: why do experiences seem to have phenomenal properties? Frankish is candid about the cost: if zombies lack only phenomenal consciousness in the philosophers’ technical sense, then “in this technical sense, we are zombies.”

Dennett’s 1988 essay “Quining qualia” set out to destroy faith in the very concept with fifteen thought experiments, and concluded that “there simply are no qualia at all”[20]. Eliminativism goes further and says that much of our everyday psychological vocabulary will be dropped by a mature science of the mind.

The objection to all of this is short. Chalmers calls it Moorean: it is more certain that people sometimes feel pain than that any premise of the illusionist argument is true[21].

Galen Strawson called the denial of experience the strangest thing in the history of human thought. Illusionists reply that an illusion of experience needn’t itself be experienced, any more than a misprint in a map needs to be printed twice.

Identity theory, functionalism and the phenomenal concept strategy

The identity theory of the 1950s said that sensations are brain processes, as lightning is an electrical discharge. Its trouble was multiple realisability: octopuses, birds and people all seem to feel pain with very different nervous systems, so pain can’t be identical to one kind of neural process[22].

Functionalism answered by identifying a mental state with its role in the system rather than with its stuff, which makes minds possible in any substrate, including silicon. It is the most popular single answer among philosophers, at 33% on the survey’s question about consciousness, ahead of dualism at 22% and identity theory at 13%.

The most common way physicalists handle the hard problem today is the phenomenal-concept strategy of Brian Loar and David Papineau. The gap between experience and brain is real, but it is a gap between two kinds of concept, not two kinds of thing: one we get by having an experience, the other from theory.

Two concepts can pick out one thing without our being able to see that they do, as “water” and “H₂O” did.

Chalmers’s objection is a dilemma: whatever feature of phenomenal concepts is supposed to explain the gap must either itself be physically explicable, in which case it can’t explain why a zombie wouldn’t have it too, or not, in which case physicalism has given up[23].

Biological naturalism and mysterianism

John Searle’s biological naturalism holds that consciousness is a real, subjective biological phenomenon, caused by the brain’s specific processes as liquidity is a feature of water, and that running the right program is not enough[24].

The objection is that it never says which causal powers of the brain matter, and if it’s their organisation, it is functionalism under another name. Anil Seth’s view that experience is rooted in our nature as living, self-regulating “beast machines” is close to it in spirit.

Colin McGinn’s mysterianism is the modest version: some natural property of the brain explains consciousness, but our minds are “cognitively closed” to it, as a dog’s mind is to calculus. “How can technicolour phenomenology arise from soggy grey matter?”[25]. It is hard to argue with and hard to make use of.

Dualism

Property dualism, Chalmers’s own view, holds that experiences are fundamental features of the world, tied to physical organisation by laws of nature we haven’t yet found, much as electromagnetism had to be added to mechanics.

Its main problem is causal: if physics is complete, non-physical experiences have nothing left to do, and then our talk about experience isn’t caused by experience.

Epiphenomenalism accepts that consequence: experience is caused by the brain and causes nothing, like the whistle of a steam train, in Thomas Huxley’s image. William James’s reply still stands: a property with no effects could not have been selected by evolution.

Substance dualism, Descartes’ view that the mind is a separate thing, has modern defenders but faces the old objection his correspondent Elisabeth of Bohemia raised in 1643: how could an unextended mind move an extended body? Today the objection is that every physical effect seems to have a sufficient physical cause, leaving nothing for a mind to do.

Panpsychism, Russellian monism and idealism

The most surprising turn in recent philosophy of mind has been the return of panpsychism, the view that some fundamental physical entities have, in Chalmers’s words, mental states[26].

Its best argument comes from Bertrand Russell and Arthur Eddington: physics describes what mass and charge do, their structure and dynamics, and is silent about what they intrinsically are. Experience is the one intrinsic nature we know from the inside, so perhaps that is what the structure is the structure of.

This is Russellian monism; Galen Strawson put the argument bluntly: “Real physicalists must accept that at least some ultimates are intrinsically experience-involving”[27]. Panpsychism does not say that rocks have feelings; a rock, like a crowd, need not be a subject.

Its own defenders name its hardest problem: combination. How do many tiny subjects add up to one big subject, mine? The existence of a hundred subjects doesn’t seem to require a hundred-and-first[28].

William James had put it in 1890: a hundred feelings packed together each remain “shut in its own skin, windowless, ignorant of what the other feelings are and mean”. On the survey, panpsychism was the answer of 8% of philosophers.

Idealism, the view that reality is fundamentally mental, survives in Bernardo Kastrup’s “analytic idealism”, where individual minds are dissociated parts of one universal consciousness. It reverses the combination problem: how does one mind split into many?

Scientific theories of consciousness

Philosophers ask what consciousness is. Neuroscientists ask which brain activity distinguishes seeing something from not seeing it, while the stimulus stays the same: a word flashed so briefly that sometimes you see it and sometimes you don’t, two images competing for one eye.

The answers have grown into theories, more than two hundred by one count[29]. Six of them matter most.

Four views of the left hemisphere. Global workspace theory: prefrontal and parietal hubs. Integrated information theory: the posterior hot zone of occipital, temporal and parietal cortex. Higher-order theory: dorsolateral prefrontal cortex. Recurrent processing theory: visual cortex by itself Four views of the left hemisphere. Global workspace theory: prefrontal and parietal hubs. Integrated information theory: the posterior hot zone of occipital, temporal and parietal cortex. Higher-order theory: dorsolateral prefrontal cortex. Recurrent processing theory: visual cortex by itself

Global neuronal workspace theory (Bernard Baars, Stanislas Dehaene, Jean-Pierre Changeux) says the brain runs many specialised processors unconsciously, and information becomes conscious when it wins entry to a limited-capacity workspace and is broadcast to all of them at once, available for report, memory and planning (Dehaene & Naccache, 2001).

Entry is a sudden, all-or-none ignition in a network of prefrontal and parietal cortex, about 300 milliseconds after a stimulus. Engineers will recognise the design: Baars took it from the “blackboard” architectures of early AI.

Integrated information theory (Giulio Tononi, Christof Koch) starts from the other end, from properties it says every experience has, and asks what a physical system must be like to have them.

Its answer: consciousness is a system’s maximally irreducible cause–effect structure, its quantity is integrated information, Φ, and its seat in the brain is a “posterior hot zone” of sensory cortex[30].

It has striking consequences: a feed-forward network has Φ = 0 however intelligent it is, and a simple grid of logic gates can have more Φ than a brain, a consequence Tononi accepted when the computer scientist Scott Aaronson pointed it out[31]. It also concedes that Φ can’t be computed for any realistic system.

Higher-order theories (David Rosenthal, Hakwan Lau, Joseph LeDoux) say a state is conscious when you are, in a suitable way, aware of being in it, through a representation of the state itself[32].

Recurrent processing theory (Victor Lamme) says visual experience begins when higher visual areas feed signals back to lower ones, no frontal cortex needed: the first fast sweep up the visual hierarchy is unconscious (Lamme & Roelfsema, 2000).

Predictive processing, in Anil Seth’s version, says perception is the brain’s best guess about the causes of its signals, a “controlled hallucination”, and that the deepest guesses are about the state of the body itself[33].

And Michael Graziano’s attention schema theory says the brain builds a simplified model of its own attention, and “the content of this internal model leads a brain to conclude that it has a subjective experience”[34], which makes it a scientific cousin of illusionism.

Testing global workspace theory against integrated information theory

The usual way theories of consciousness are tested makes it easy for each to declare victory, since each team runs the experiments its theory suggests.

In 2018 the Templeton World Charity Foundation brought proponents of the two leading theories, Tononi and Koch for integrated information and Dehaene for the global workspace, together to do something different: agree in advance on three predictions where the theories conflict and on what result would count against each, and let theory-neutral labs collect the data[35].

Participants looked at clearly visible faces, objects, letters and false fonts for half a second to a second and a half while 256 of them were recorded by fMRI, MEG or electrodes in the brains of epilepsy patients[2].

The adversarial test of integrated information theory and global neuronal workspace theory. First prediction, where conscious content lives: IIT predicted posterior cortex is enough, GNWT that prefrontal cortex is needed; content was decodable at the back and adding prefrontal cortex added nothing, a pass for IIT and mixed for GNWT. Second, how it lasts: IIT predicted sustained posterior activity, GNWT ignition at onset and offset; duration was tracked, sparsely, a pass for IIT, while none of 655 prefrontal electrodes showed offset ignition, a failure for GNWT. Third, connectivity: IIT predicted sustained posterior synchrony, which was absent, a failure; GNWT brief long-range synchrony, not supported by the preregistered metric but consistent in an exploratory analysis, mixed The adversarial test of integrated information theory and global neuronal workspace theory. First prediction, where conscious content lives: IIT predicted posterior cortex is enough, GNWT that prefrontal cortex is needed; content was decodable at the back and adding prefrontal cortex added nothing, a pass for IIT and mixed for GNWT. Second, how it lasts: IIT predicted sustained posterior activity, GNWT ignition at onset and offset; duration was tracked, sparsely, a pass for IIT, while none of 655 prefrontal electrodes showed offset ignition, a failure for GNWT. Third, connectivity: IIT predicted sustained posterior synchrony, which was absent, a failure; GNWT brief long-range synchrony, not supported by the preregistered metric but consistent in an exploratory analysis, mixed

Both theories passed something and failed something. For integrated information, the content of experience could be decoded from posterior cortex, and adding prefrontal cortex didn’t help, as predicted; but the sustained synchronisation within posterior cortex that the theory predicted was absent.

For the global workspace, the category of the object was decodable in prefrontal cortex, but the theory predicted a second ignition when the stimulus disappeared, and “none of the 655 electrodes measured the temporal profile predicted by GNWT”, while the identity and orientation of faces, both clearly seen, were barely represented in prefrontal cortex at all.

The consortium’s summary: the results “align with some predictions of IIT and GNWT, while substantially challenging key tenets of both theories.”

Two lessons came out of it, and neither is the one the headlines wanted. First, the experiment tested where and when each theory expects activity in the brain, not its core definition of consciousness; as the authors wrote, theorists “could respond to challenging data by modifying the proposed biological implementation while retaining the mathematical or computational core”.

Second, the early press coverage, which presented the results as a win for integrated information, provoked 124 researchers to sign a letter calling the theory pseudoscience, and a public dispute that ended up in Nature Neuroscience[36, 37].

A quieter result may matter more. A brain wave called the P3b, appearing about 300 ms after a stimulus, was long treated as a signature of consciousness.

When people saw stimuli but didn’t have to report them, it vanished: “we found a robust P3b in the report condition, but no P3b whatsoever in the no-report condition”[38]. Much of what looked like the brain becoming conscious of something may have been the brain preparing to say so.

Covert consciousness in unresponsive patients

The question stops being academic in an intensive care unit. In 2006 Adrian Owen and colleagues put a young woman who met the clinical criteria for a vegetative state into a scanner and asked her to imagine playing tennis, then to imagine walking through her home.

Her brain lit up in the motor and spatial areas that light up in healthy volunteers, “in a manner indistinguishable from that of healthy volunteers”[39].

In 2024 a study at six centres tested 353 adults with disorders of consciousness. Of the 241 who showed no observable response to commands, 60, a quarter, performed the imagery tasks in fMRI or EEG[40].

The group included patients in coma and in a minimally conscious state, not only vegetative ones, and the tests missed many patients who could respond. But a quarter is a lot of people misjudged.

Two grids of squares. Of 241 patients with no response to commands at the bedside, 60 performed the mental imagery tasks on fMRI or EEG. Of 112 patients who did follow commands at the bedside, only 43 showed it on the scanner tests Two grids of squares. Of 241 patients with no response to commands at the bedside, 60 performed the mental imagery tasks on fMRI or EEG. Of 112 patients who did follow commands at the bedside, only 43 showed it on the scanner tests

A test inspired by integrated information theory, the perturbational complexity index, takes a different route. It knocks on the cortex with a magnetic pulse and measures how complex the echo is in the EEG: a conscious brain responds with activity that is both widespread and varied, an unconscious one with a simple, local or uniform response[41].

In a benchmark of 150 subjects, awake, dreaming, deeply asleep and under different anaesthetics, a single threshold of 0.31 separated conscious from unconscious states without a single error, and it caught 36 of 38 minimally conscious patients[42].

Nine of 43 patients diagnosed as vegetative had high complexity, a hint that some of them might be conscious and unable to show it. The index measures a capacity, calibrated against what people later report, and its success doesn’t make integrated information theory true.

Can AI be conscious?

Every theory above gives a different answer for a language model, which is why the question can’t be settled by asking the model.

Butlin, Long and seventeen colleagues took the theories that are compatible with computational functionalism, the assumption that the right computation is enough, and turned each into “indicator properties”: recurrent processing, a global workspace, higher-order monitoring, an attention schema, predictive processing, agency and embodiment.

They checked current systems against the list and concluded: “Our analysis suggests that no current AI systems are conscious,” adding that “there are no obvious technical barriers to building AI systems which satisfy these indicators”[43].

Chalmers put numbers on his uncertainty, explicitly as illustrations: under 10% for current language models, and 25% or more that within a decade there will be extended systems, with senses, memory, a workspace and agency, that are conscious[44].

Integrated information theory gives the opposite answer for a different reason: a conventional computer would “experience next to nothing” whatever it computes, because what matters is the causal structure of the hardware, not the software[45].

Seth argues that consciousness may depend on being alive, so “real artificial consciousness is unlikely along current trajectories”[46]; in October 2026 he, Block and Tononi wrote together that they “feel that consciousness in AI is extraordinarily unlikely” in today’s digital systems[47].

Meanwhile the functional side keeps moving. In July 2026 interpretability researchers at Anthropic reported that language models “maintain a small, privileged set of representations bearing some of the functional hallmarks of conscious access”, something like a global workspace[48].

That is a claim about access, Block’s first kind of consciousness, not about experience, and it is exactly the kind of finding that each position will read differently: as a step toward machine consciousness, as irrelevant to it, or as a model of how a system comes to believe it is conscious.

The last reading is the one Chalmers himself recommends as common ground. The meta-problem of consciousness is “the problem of explaining why we think that there is a problem of consciousness”[21].

It is an easy problem, a question about behaviour and mechanisms, and it can now be studied in machines: why do systems trained on human text say they have an inner life, and what in them produces the saying?

Chalmers wrote that if he were a materialist, he would be an illusionist. An illusionist would say that answering the meta-problem is all there ever was to answer.

What a physicalist can conclude

Five conclusions seem fair.

The hard problem is not a gap in the data. It is the observation that explaining every function of the brain, which is the job of neuroscience, leaves open why the functions are felt.

You can deny that this is a further question, as type-A physicalists do, but then denying it is your theory of consciousness, and you should be ready to defend it.

“Physicalism” is a decision, not a theory. Within it you still have to choose whether consciousness is its function or its biology, whether the hard problem is real and how to deal with it, and the choice decides what you think about machines.

The thought experiments are tools for finding your position, not proofs. Mary’s room converted its own author. The fading-qualia argument for silicon consciousness came from a dualist.

The science is progressing, mostly by eliminating things. Report is not experience, prefrontal cortex is not obviously where content lives, and no current theory survived its first preregistered test intact.

The practical questions, which patients are conscious and whether a machine ever could be, won’t wait for the philosophical one to be settled. The best we can do is to be clear about which theory each answer assumes.

The figures as code: articles/consciousness.

References

  1. Lenharo, 2023. doi.org. ↩
  2. Cogitate Consortium, 2025. doi.org. ↩
  3. Nagel, 1974. doi.org. ↩
  4. Chalmers, 1995. consc.net. ↩
  5. Levine, 1983. informationphilosopher.com. ↩
  6. Block, 1995. doi.org. ↩
  7. Smart, 1959. doi.org. ↩
  8. Chalmers, 2003. consc.net. ↩
  9. Bourget & Chalmers, 2023. doi.org. ↩
  10. Jackson, 1982. doi.org. ↩
  11. Jackson, 1986. doi.org. ↩
  12. Loar, 1990. doi.org. ↩
  13. Nida-Rümelin & O Conaill, 2024. plato.stanford.edu. ↩
  14. Chalmers, 1996. global.oup.com. ↩
  15. Kirk, 2023. plato.stanford.edu. ↩
  16. Kripke, 1980. hup.harvard.edu. ↩
  17. Byrne, 2025. plato.stanford.edu. ↩
  18. Block, 1980. doi.org. ↩
  19. Frankish, 2016. keithfrankish.com. ↩
  20. Dennett, 1988. doi.org. ↩
  21. Chalmers, 2018. consc.net. ↩
  22. Putnam, 1967. plato.stanford.edu. ↩
  23. Chalmers, 2007. consc.net. ↩
  24. Searle, 1992. doi.org. ↩
  25. McGinn, 1989. doi.org. ↩
  26. Chalmers, 2015. consc.net. ↩
  27. Strawson, 2008. doi.org. ↩
  28. Chalmers, 2017. doi.org. ↩
  29. Kuhn, 2024. doi.org. ↩
  30. Albantakis et al., 2023. doi.org. ↩
  31. Aaronson, 2014. scottaaronson.blog. ↩
  32. Brown et al., 2019. doi.org. ↩
  33. Seth & Tsakiris, 2018. doi.org. ↩
  34. Graziano & Webb, 2015. doi.org. ↩
  35. Melloni et al., 2023. doi.org. ↩
  36. IIT-Concerned et al., 2025. doi.org. ↩
  37. Tononi et al., 2025. doi.org. ↩
  38. Cohen et al., 2020. doi.org. ↩
  39. Owen et al., 2006. doi.org. ↩
  40. Bodien et al., 2024. doi.org. ↩
  41. Casali et al., 2013. doi.org. ↩
  42. Casarotto et al., 2016. doi.org. ↩
  43. Butlin et al., 2023. arxiv.org. ↩
  44. Chalmers, 2023. arxiv.org. ↩
  45. Tononi & Koch, 2015. doi.org. ↩
  46. Seth, 2025. doi.org. ↩
  47. Seth et al., 2026. doi.org. ↩
  48. Gurnee et al., 2026. arxiv.org. ↩
]]> Eva Koroleva How Karl Friston’s free energy principle works and what its critics say 2026-06-23T00:00:00+00:00 2026-11-08T00:00:00+00:00 https://blog.korolevalab.com/blog/2026/free-energy-principle

In short

The free energy principle says that anything that keeps existing can be described as minimising one quantity, variational free energy, an upper bound on how surprising its sensations are under its own model of the world.

It does so in three ways: by changing its beliefs (perception), its model (learning) and the world (action). The mathematics is the same as the objective machine learning maximises in variational autoencoders, with the sign flipped.

As a principle it can’t be falsified, and its author says so; what can be tested are the specific theories built on it, like predictive coding, and the evidence for those is real but mixed.

≈8,000citations of Friston's 2010 paper, depending on the database
<4%error predicting how rat neurons in a dish would rewire, from their first 10 sessions
"mixed"the verdict of a 2020 review on the brain evidence for predictive processing

When Karl Friston was eight, he watched woodlice in a hot garden. They scurried about, and he decided that the scurrying had no purpose: they simply moved faster in the sun and slower in the shade.

Nobody told a woodlouse to seek the shade; it ended up there because it slowed down when it got there. He has said that the same principle underwrites the mathematics he spent the next half century developing[1].

Friston is best known to neuroscientists for something more practical: statistical parametric mapping, the method behind a large share of all brain-imaging studies[2].

In 2016 a ranking of influence in neuroscience by Semantic Scholar put him first[3]. His 2010 paper “The free-energy principle: a unified brain theory?” has been cited about eight thousand times[4]. And the principle is also famous for something else: almost nobody can explain it.

This post tries. It uses the formulas, because the principle is a mathematical statement and loses its meaning without them, but every symbol is explained, and each step comes with what it means in plain words. It ends with what the principle predicts, what has been tested, and the strongest arguments against it.

What is the free energy principle?

Anything that persists, a cell, a fish, a person, stays within a small range of states: a fish stays in water, your body temperature stays near 37 °C.

In the language of probability, it keeps its sensations unsurprising. It can’t measure surprise directly, so it does something cleverer: it keeps a model of the world, and it minimises a quantity called variational free energy, which can be computed from that model and is never smaller than surprise.

There are three ways to make free energy smaller: change your beliefs to fit what you sense, which is perception; change your model over time, which is learning; or change what you sense by acting on the world, which is action. That is the whole principle; the rest is working out what it implies.

Surprisal and why living things keep it low

Start with a definition. If a creature’s model of the world assigns probability \(p(o)\) to an observation \(o\), the surprisal of that observation is

\[\Im(o) = -\ln p(o).\]

An observation the model expects, with probability near 1, has surprisal near 0; an improbable one has large surprisal. Surprisal is a property of the observation relative to a model, not a feeling.

A fish’s model of its world assigns high probability to being wet, so being on dry land is highly surprising, in this technical sense and in every other.

Friston’s argument for why living things should keep surprisal low runs as follows. Anything that stays the kind of thing it is spends its time in a limited set of states. The probability distribution over its states is therefore concentrated, which means it has low entropy.

Entropy is average surprisal, and if the system is ergodic, if averaging over time gives the same answer as averaging over all the states it could be in, then keeping surprisal low moment to moment keeps entropy low over a lifetime[5].

“Biological agents must therefore minimize the long-term average of surprise to ensure that their sensory entropy remains low”[4].

Two traps

This entropy is the Shannon entropy of a probability distribution, not the entropy of thermodynamics; living things don’t violate the second law, they keep their own states ordered while exporting disorder to their surroundings. And ergodicity is an assumption, not a fact about organisms, which do grow, age and die.

Variational free energy and the ELBO

There is a problem. To compute \(p(o)\) a creature would have to add up the probability of the observation over every possible cause of it, every hidden state \(x\) of the world that could have produced it:

\[p(o) = \int p(o \mid x)\, p(x)\, dx.\]

That integral is intractable for anything interesting. As Friston puts it, “a system cannot know whether its sensations are surprising and could not avoid them even if it did know”[4].

The way out is an old trick from statistics. The creature keeps a belief about the causes, a probability distribution \(q(x)\) that it can represent, for instance in neural activity, and scores that belief with the variational free energy:

\[F[q, o] = \mathbb{E}_{q(x)}\big[\ln q(x) - \ln p(o, x)\big].\]

Here \(p(o, x) = p(o \mid x)\,p(x)\) is the creature’s generative model: how likely each hidden cause is, and what observations each cause would produce. \(F\) needs only the model and the belief, both of which the creature has. One line of algebra rewrites it in the form that explains everything else:

\[F = D_{\mathrm{KL}}\big[q(x) \,\Vert\, p(x \mid o)\big] - \ln p(o).\]

The first term is the Kullback–Leibler divergence between the belief and the true posterior, the distribution over causes that perfect Bayesian inference would give; it is never negative and is zero only when the belief is exactly right.

The second term is the surprisal. So free energy is surprisal plus a non-negative error: \(F \ge \Im(o)\), whatever the belief. Lower \(F\) by improving the belief, and two things happen at once: the belief gets closer to the truth, and \(F\) becomes a tighter estimate of surprisal.

Free energy against the centre of a Gaussian belief about the size of the food. A horizontal line marks the surprisal of the observation, about 3.24 nats. Free energy for beliefs of a fixed width and for the best width at each centre always stays above the line; the best belief, centred near 1.5, comes within a few hundredths of a nat of it Free energy against the centre of a Gaussian belief about the size of the food. A horizontal line marks the surprisal of the observation, about 3.24 nats. Free energy for beliefs of a fixed width and for the best width at each centre always stays above the line; the best belief, centred near 1.5, comes within a few hundredths of a nat of it

The chart computes this for the example of the next section: free energy for a range of beliefs, each a bell curve centred at a different value, against the surprisal of the observation. No belief gets below the line; the best ones nearly touch it.

The same quantity has a second decomposition, which explains why the principle has so much to say about learning:

\[F = \underbrace{D_{\mathrm{KL}}\big[q(x) \,\Vert\, p(x)\big]}_{\text{complexity}} - \underbrace{\mathbb{E}_{q(x)}\big[\ln p(o \mid x)\big]}_{\text{accuracy}}.\]

Accuracy is how well the believed causes predict the observation. Complexity is how far the belief has moved from what the creature expected before it looked. Minimising free energy is finding the simplest explanation that fits the data, Occam’s razor written as an equation.

Three ways to write the same variational free energy. First: the divergence between the belief and the true posterior, which is never negative, plus the surprisal of the observation, so free energy is an upper bound on surprisal. Second: complexity, how far the belief moved from the prior, minus accuracy, how well the believed causes predict the observation. Third: the negative of the evidence lower bound that variational autoencoders maximise Three ways to write the same variational free energy. First: the divergence between the belief and the true posterior, which is never negative, plus the surprisal of the observation, so free energy is an upper bound on surprisal. Second: complexity, how far the belief moved from the prior, minus accuracy, how well the believed causes predict the observation. Third: the negative of the evidence lower bound that variational autoencoders maximise

If you work in machine learning, you have met \(F\) already. A variational autoencoder is trained to maximise the evidence lower bound, the ELBO, which is exactly \(-F\) with the data called \(x\) and the hidden causes called \(z\)[6].

Active-inference papers say it outright: “The variational free energy is the negative evidence lower bound that is optimised in variational Bayes in machine learning”[7].

The idea of a brain doing this goes back to the Helmholtz machine of Dayan, Hinton, Neal and Zemel in 1995[8], named after Hermann von Helmholtz, who described perception as unconscious inference in the 1860s.

Why the name

The formula has the same shape as the Helmholtz free energy of statistical physics, energy minus temperature times entropy, if you call \(-\ln p(o, x)\) the energy and set the temperature to 1. The mathematics is shared; there are no joules in it.

Predictive coding and precision

Free energy says what perception optimises, not how a brain could do it. To get a mechanism, Friston assumes that the generative model is Gaussian: an observation is a prediction from the hidden cause plus noise, and the cause has a bell-shaped prior.

Then free energy turns, up to constants, into a sum of squared prediction errors, each weighted by its precision, the inverse of its variance[9]:

\[\begin{aligned} F(\mu) &\approx \tfrac12\,\Pi_o\,\big(o - g(\mu)\big)^2 \\ &\quad {}+ \tfrac12\,\Pi_x\,\big(\mu - \eta\big)^2 + \text{const}. \end{aligned}\]

Here \(\mu\) is the creature’s best guess about the hidden cause, \(g(\mu)\) is the observation that guess predicts, \(\eta\) is what it expected before looking, and \(\Pi_o\) and \(\Pi_x\) are how much it trusts its senses and its prior.

Perception is then gradient descent: the guess moves until the error from below, the senses, balances the error from above, the prior.

Rafal Bogacz’s tutorial makes this concrete with a creature that has one light receptor and wants to know how big a piece of food is. Light grows with the area of the food, so a food of size \(v\) gives light \(g(v) = v^2\) plus noise.

The creature expects food of size 3, observes a light intensity of 2, and both the prior and the sensor have variance 1. A size of 3 would give light 9, far from 2; a size of \(\sqrt 2 \approx 1.41\) would match the light perfectly but contradicts the prior. Where does the guess settle?

Left: the exact posterior over the food's size, peaked at 1.57 with a mean of 1.48. Right: the creature's guess over time, starting from its prior expectation of 3. Following the gradient of free energy it settles at 1.57 within about one time unit; a network with explicit prediction-error units oscillates and settles at the same value later Left: the exact posterior over the food's size, peaked at 1.57 with a mean of 1.48. Right: the creature's guess over time, starting from its prior expectation of 3. Following the gradient of free energy it settles at 1.57 within about one time unit; a network with explicit prediction-error units oscillates and settles at the same value later

The left panel is the exact answer, computed by Bayes’ rule on a grid. The right panel is the creature’s guess following the gradient of free energy from its prior expectation, computed here exactly as in Bogacz’s exercises.

It settles near 1.57, at the peak of the exact posterior. Bogacz’s text: “it rapidly converges to the value of ϕ≈1.6, which is also the value that maximizes the exact posterior probability.”

The guess finds the most probable size, not the average one, which is about 1.48 for this lopsided posterior; a single number cannot carry uncertainty, which is one of the approximations the scheme pays for its simplicity.

How the charts in this post are computed

They come from the formulas above, not from the papers’ figures. For the bound, the joint probability of Bogacz’s example, a normal prior on the size with mean 3 and variance 1 times a normal likelihood of the light with mean \(v^2\) and variance 1, is evaluated on a fine grid of sizes.

Surprisal is minus the log of its integral, and free energy is the integral of \(q(v)\,[\ln q(v) - \ln p(u, v)]\) for each Gaussian belief \(q\).

The guess follows Bogacz’s second exercise: small Euler steps along the gradient of \(\ln p(v) + \ln p(u \mid v)\), starting at the prior mean. The thinner line is his third exercise, where the same descent runs through two prediction-error units, one for the prior and one for the senses. The risk and ambiguity bars further down use the two-state examples of Smith, Friston and Whyte’s tutorial, combined by the formula for \(G\).

Put many such units in a hierarchy and you get predictive coding. Each level sends predictions down to the level below and receives back the error: what the prediction failed to explain. Rao and Ballard proposed in 1999 that feedback connections in the visual cortex “carry predictions of lower-level neural activities” and feedforward connections carry the residual errors[10].

A predictive-coding hierarchy of three levels above the senses. Each level sends a prediction down to the level below and receives back the prediction error, what the prediction failed to explain, weighted by its precision. Precision works like a volume knob on each error: turned up, the guess follows the data; turned down, it follows expectations. In Friston's account attention is precision A predictive-coding hierarchy of three levels above the senses. Each level sends a prediction down to the level below and receives back the prediction error, what the prediction failed to explain, weighted by its precision. Precision works like a volume knob on each error: turned up, the guess follows the data; turned down, it follows expectations. In Friston's account attention is precision

Precision is the most interesting knob in the scheme. It sets how loudly each error is heard.

Turn up the precision of the senses and the guess follows the data; turn up the precision of the prior and the guess follows expectations. In Friston’s account attention is precision: attending to something means trusting its prediction errors more.

With Gaussian assumptions, minimising free energy is gradient descent on prediction errors weighted by precision. It finds the most probable cause, not a full posterior.

Active inference and expected free energy

Here the free energy principle departs from ordinary Bayesian perception. Free energy depends on the observation, and a creature can change its observations by acting. If your model predicts that your arm is raised, there are two ways to remove the error: revise the belief, or raise the arm.

Friston calls the second active inference, and in his account movement works this way: the motor system receives predictions of how the body should feel, and reflexes fulfil them[11].

Because only the accuracy term depends on what is sensed, “action can reduce free energy only by increasing accuracy (that is, selectively sampling data that are predicted)”[4].

To plan, a creature needs to score futures it hasn’t seen yet.

For that, active inference uses a second quantity, the expected free energy of a policy \(\pi\), a sequence of actions[12]. It too can be split in two ways. One reads as risk plus ambiguity[13]:

\[\begin{aligned} G(\pi) &= \underbrace{D_{\mathrm{KL}}\big[q(o \mid \pi) \,\Vert\, p(o \mid C)\big]}_{\text{risk}} \\ &\quad {}+ \underbrace{\mathbb{E}_{q(x \mid \pi)}\big[H[p(o \mid x)]\big]}_{\text{ambiguity}}. \end{aligned}\]

Risk is how far the outcomes the policy is expected to bring are from the outcomes the creature prefers, \(p(o \mid C)\). Ambiguity is how noisy the link from states to observations will be: a policy that leads somewhere you can’t see clearly is ambiguous.

The other split reads the same quantity as minus the expected information gain, how much the outcomes would teach about the world, minus the pragmatic value, how well they match preferences. A creature choosing policies with low expected free energy is curious and goal-directed by the same token.

Two bar charts. Risk: a policy expected to reach the preferred outcome 82% of the time has risk 2.41 nats, a policy that is a coin flip 7.31. Ambiguity: a policy heading toward the murky state has ambiguity 0.66 nats, one heading toward the clearly observed state 0.52. Lower expected free energy is preferred Two bar charts. Risk: a policy expected to reach the preferred outcome 82% of the time has risk 2.41 nats, a policy that is a coin flip 7.31. Ambiguity: a policy heading toward the murky state has ambiguity 0.66 nats, one heading toward the clearly observed state 0.52. Lower expected free energy is preferred

One feature of this needs saying plainly, because it surprises everyone. There is no reward.

Preferences enter as a probability distribution over observations: the creature expects to see what it prefers, and acts to make those expectations come true. Friston’s colleagues call this an optimism bias written into the model.

It also answers the oldest objection to the principle, the dark room problem. If creatures minimise surprise, why don’t they find a dark, quiet room and stay there?

Friston’s reply is that surprise is always relative to the creature’s model, and a creature whose model, shaped by evolution, predicts a varied world full of food and company finds a dark room surprising[14].

Expected free energy adds a second answer: a dark room is maximally ambiguous, so the first thing an active-inference agent does there is switch on the light.

Markov blankets

The principle’s most ambitious version starts from physics, not brains. It asks: what does it take for anything to exist as something separate from its surroundings?

In probability, a Markov blanket of a variable is a set of other variables that, once known, makes everything else irrelevant to it. In a Bayesian network it is the variable’s parents, children and its children’s other parents, a notion due to Judea Pearl.

Friston carries it into dynamics: a thing is a set of internal states separated from external states by a blanket of sensory states, which the outside affects and the inside doesn’t, and active states, which the inside affects and the outside doesn’t[15].

A cell’s membrane is the obvious example; a candle flame, Friston notes, isn’t one, because its molecules are swept away almost instantly.

A Markov blanket. External states act on sensory states; sensory states act on internal states; internal states act on active states; active states act on external states. Sensory and active states form the blanket: given the blanket, internal and external states are independent. A cell membrane is an example; a candle flame, whose molecules are swept away almost instantly, is not A Markov blanket. External states act on sensory states; sensory states act on internal states; internal states act on active states; active states act on external states. Sensory and active states form the blanket: given the blanket, internal and external states are independent. A cell membrane is an example; a candle flame, whose molecules are swept away almost instantly, is not

The 2019 monograph “A free energy principle for a particular physics” claims that for any system that settles into a characteristic set of states and has such a blanket, the internal states can be read as if they were performing Bayesian inference about the external ones, minimising free energy[16].

The assumptions are many and technical: random dynamics of a particular form, a steady state that the system keeps returning to, coupling sparse enough to produce a blanket, and restrictions on how internal and external states can circulate around each other. Several of them have been challenged, as we’ll see.

Experimental evidence

The principle itself makes no prediction that could fail; we’ll come back to that. The process theories built on it do.

Predictive coding in the brain. Expectation effects in the brain are well documented: the mismatch negativity, a brain response to an unexpected sound in a sequence, is robust and widely replicated. Whether the brain implements predictive coding in particular is less clear.

A 2020 review, co-written by Andy Clark, one of the framework’s best-known philosophers, concluded that “on the whole it must be concluded that the empirical support for each of PP’s key hypotheses is mixed”[17]. A 2024 audit by an active-inference modeller found “modest support”[18].

Recordings in monkey visual cortex found responses to expected and pattern-violating stimuli “nearly identical”[19], while neurons in layer 2/3 of mouse visual cortex compute a difference between what the animal’s movement predicts and what it sees, as a prediction-error unit should[20].

Neurons in a dish. The most direct test of the principle itself came in 2023. Isomura, Friston and colleagues grew rat cortical neurons on a grid of 64 electrodes and stimulated them with mixtures of two hidden signals.

The networks learned to separate the sources, and a model derived from free-energy minimisation, set up from the first 10 training sessions alone, predicted how the networks’ connections would change over the next 90 with an error under 4%[21].

The authors are careful about what that shows: the principle “is sometimes considered to be experimentally irrefutable”, so they framed this as a test of predictive power, and “showing the existence of such a mapping alone is insufficient as an empirical validation.”

Pong. A year earlier, the company Cortical Labs and Friston’s group taught neurons in a dish to play the video game Pong. Missing the ball triggered unpredictable stimulation; the cultures learned to rally longer within five minutes, which the authors explained as acting to avoid unpredictable input[22].

The paper’s title said the neurons “exhibit sentience”, in the minimal sense of being “responsive to sensory impressions”, and thirty neuroscientists wrote to the journal to object[23].

A detail matters more than the word: cultures that got no stimulation at all after a miss also learned, less well, which weakens the reading that unpredictability was what taught them.

Psychiatry. The idea that some symptoms reflect aberrant precision, trusting the senses too little or too much relative to expectations, has been applied to psychosis[24] and autism[25].

Adults with autism were found to be “less surprised than neurotypical adults when their expectations are violated” in a volatile task[26]. For psychosis the direction is disputed: “some studies implicate weakened prior beliefs in psychosis, and others find stronger priors”[27].

A scorecard of claims. That free energy bounds surprisal is a derivation and true. That brains signal violations of expectation is supported by the mismatch negativity and by mismatch neurons in mouse visual cortex. That cortex runs hierarchical predictive coding has mixed or modest support. That neurons in a dish minimise free energy is supported in one preparation. Aberrant precision in psychosis and autism: mixed. That active inference beats reinforcement learning in AI: weak. That every persisting system minimises free energy, the 2019 physics: contested, holding in narrow cases. The principle itself: not testable, by design A scorecard of claims. That free energy bounds surprisal is a derivation and true. That brains signal violations of expectation is supported by the mismatch negativity and by mismatch neurons in mouse visual cortex. That cortex runs hierarchical predictive coding has mixed or modest support. That neurons in a dish minimise free energy is supported in one preparation. Aberrant precision in psychosis and autism: mixed. That active inference beats reinforcement learning in AI: weak. That every persisting system minimises free energy, the 2019 physics: contested, holding in narrow cases. The principle itself: not testable, by design

The brain clearly uses expectations. Whether it computes them by predictive coding is mixed in the evidence, and direct tests of the principle itself are still few.

Criticism of the free energy principle

It can’t be falsified, and its author agrees. “The free energy principle is what it is—a principle. Like Hamilton’s Principle of Stationary Action, it cannot be falsified”[1].

Friston’s colleagues defend this as the difference between a principle, “a piece of mathematical reasoning, which can be developed into a method”, and the theories built with it, which can be tested once mapped to data; they also concede that “the FEP itself has nothing instructive to say of the demarcations defining life or consciousness per se”[28].

Critics accept the distinction and draw the consequence: if the principle explains any behaviour, it explains none in particular, and the work of explanation is done by the specific models.

Samuel Gershman put it most precisely: “FEP does not have a fixed set of distinctive claims. Rather, it makes different claims under different sets of assumptions”[29].

The Markov blankets are doing a different job from Pearl’s. Bruineberg, Dołęga, Dewhurst and Baltieri distinguish “Pearl blankets”, a tool for reasoning about a model, from “Friston blankets”, a claim about where things end in the world, and argue that moving from one to the other “requires additional philosophical premises and cannot be justified by an appeal to the success of the mathematical framework alone”[30].

The physics holds only in special cases. Aguilera, Millidge, Tschantz and Buckley checked the 2019 derivations on the simplest systems where they can be checked, linear stochastic ones, and found the required conditions “are only valid for a very narrow space of parameters”; a central step, equating the dynamics of average states with the average of the dynamics, “does not hold in general even for linear stochastic systems”[31].

Biehl, Pollock and Kanai showed by counterexample that the original “free energy lemma, when taken at face value, is wrong”, while noting that this doesn’t settle whether the general idea is worth pursuing[32].

Curiosity is put in, not derived. Millidge, Tschantz and Buckley showed that the natural extension of free energy into the future discourages exploration rather than encouraging it: “exploration does not directly follow from free energy minimization into the future”[33]. Expected free energy is a choice, not a consequence.

In AI, it hasn’t yet won anything. Comparisons of active-inference agents with reinforcement learning are mostly on small problems, and one careful study found that among deep active-inference models, “the models able to solve the dSprites environment are the ones that maximize rewards”[34].

The machine-learning connection runs the other way: predictive coding networks, it turns out, can approximate backpropagation, “converges asymptotically (and in practice, rapidly) to exact backprop gradients on arbitrary computation graphs”[35].

How to read the free energy principle

The free energy principle is best read as three things stacked on each other, and most arguments about it come from mixing them up.

At the bottom is mathematics: variational free energy is a bound on surprisal, and minimising it is approximate Bayesian inference. That part is uncontroversial and is used every day in machine learning under another name.

In the middle is a framework: perception, learning and action described as minimising the same quantity, with attention as precision and preferences as expectations. It is a genuinely useful way to think, and it has produced testable process theories, predictive coding first among them. Those theories have some support and some trouble, as theories should.

At the top are the grand claims: that every living thing, by virtue of existing, can be described as inferring its world, and that this tells us what life, mind or consciousness are.

Here the principle’s defenders themselves concede that it says little about life or consciousness as such, and its critics argue that the physics behind it holds only in special cases.

If you take one thing from it, take the middle layer: the brain as a machine that predicts, weighs its errors by how much it trusts them, and acts to make its predictions come true. Unlike the principle, that picture makes predictions an experiment can refute.

The calculations and the figures as code: articles/friston.

References

  1. Fortier & Friedman, 2018. aliusresearch.org. ↩
  2. Friston et al., 1994. doi.org. ↩
  3. UCL, 2016. web.archive.org. ↩
  4. Friston, 2010. doi.org. ↩
  5. Friston, 2012. doi.org. ↩
  6. Kingma & Welling, 2014. arxiv.org. ↩
  7. Da Costa et al., 2020. doi.org. ↩
  8. Dayan et al., 1995. doi.org. ↩
  9. Bogacz, 2017. doi.org. ↩
  10. Rao & Ballard, 1999. doi.org. ↩
  11. Friston, Kilner & Harrison, 2006. doi.org. ↩
  12. Friston et al., 2017. doi.org. ↩
  13. Smith et al., 2022. doi.org. ↩
  14. Friston, Thornton & Clark, 2012. doi.org. ↩
  15. Friston, 2013. doi.org. ↩
  16. Friston, 2019. arxiv.org. ↩
  17. Walsh et al., 2020. doi.org. ↩
  18. Hodson et al., 2024. doi.org. ↩
  19. Solomon et al., 2021. doi.org. ↩
  20. Jordan & Keller, 2020. doi.org. ↩
  21. Isomura et al., 2023. doi.org. ↩
  22. Kagan et al., 2022. doi.org. ↩
  23. Balci et al., 2023. doi.org. ↩
  24. Adams et al., 2013. doi.org. ↩
  25. Lawson et al., 2014. doi.org. ↩
  26. Lawson et al., 2017. doi.org. ↩
  27. Sterzer et al., 2018. doi.org. ↩
  28. Ramstead et al., 2023. doi.org. ↩
  29. Gershman, 2019. doi.org. ↩
  30. Bruineberg et al., 2022. doi.org. ↩
  31. Aguilera et al., 2022. doi.org. ↩
  32. Biehl et al., 2021. doi.org. ↩
  33. Millidge et al., 2021. doi.org. ↩
  34. Champion et al., 2024. doi.org. ↩
  35. Millidge et al., 2022. doi.org. ↩
]]> Eva Koroleva How similar is the human brain to AI models? 2026-06-08T00:00:00+00:00 2026-11-08T00:00:00+00:00 https://blog.korolevalab.com/blog/2026/brain-and-ai-models

In short

Neural networks trained only to do a task, never fitted to brains, turn out to predict the brain’s responses better than any model neuroscientists built by hand: language models for the language network, vision models for the visual cortex, a reinforcement-learning rule for dopamine.

The match is real and also narrower than the headlines. Better models stopped becoming more brain-like years ago, much of the measured similarity comes from word meaning and from loose statistics, and a language model is a model of one network in the brain, not of a mind.

The brain wins on energy and data by orders of magnitude, the models win on speed and breadth, and the largest gaps are lifelong learning, a body and knowing what one doesn’t know.

−97.5%drop in the language network's response to sentences a GPT-based model chose to silence it
10⁵–10⁶×more text in Llama 3.1's training than the words a child hears by adolescence
57.5%for the best video model on an intuitive-physics test where people score 96.4%

In 2023 a team at MIT trained a model on one thing: how strongly the language areas of five people’s brains responded, in a scanner, to a thousand ordinary sentences.

The model read each sentence through GPT-2 and predicted the response from GPT-2’s internal activity. Then the researchers turned it around. They let the model search 1.8 million sentences for the ones it expected to drive the language network as hard as possible, and the ones it expected to silence it.

The winners to drive were odd: “People on Insta Be Like, “Gross!”” and “Turin loves me not, nor will.” The winners to silence were the dullest sentences imaginable: “We were sitting on the couch.”

Played to new people the model had never seen, the drive sentences raised the response of the language network by 85.7% over ordinary sentences and the suppress sentences lowered it by 97.5%, on a standardised scale[1].

A language model had learned enough about one part of the human brain to steer it with the choice of words alone.

Bars of the language network's response in new people. Sentences a GPT-based encoding model chose to drive it, such as People on Insta Be Like, Gross, and Turin loves me not, nor will, raised it by 85.7% over ordinary sentences. Sentences chosen to silence it, such as We were sitting on the couch, lowered it by 97.5%, on a standardised scale Bars of the language network's response in new people. Sentences a GPT-based encoding model chose to drive it, such as People on Insta Be Like, Gross, and Turin loves me not, nor will, raised it by 85.7% over ordinary sentences. Sentences chosen to silence it, such as We were sitting on the couch, lowered it by 97.5%, on a standardised scale

That experiment is the best short answer to the question everyone asks about AI and the brain: are they alike? In some places they are alike enough to be useful to each other.

In others the resemblance is an artefact of how we measure it, and in others there is nothing to compare, because one side has no counterpart at all. This post goes through the evidence system by system, language, vision, space and memory, reward, bodies, then asks which is more efficient and what is missing.

How do you compare a brain with a model?

A brain and a neural network don’t share units. A voxel in an fMRI scan averages the blood flow around hundreds of thousands of neurons over a couple of seconds; a unit in a network is one number. So nobody compares them directly.

The standard method is an encoding model: show the same stimuli, a story, a set of photographs, to a person in a scanner and to a network, record the network’s internal activity, and fit a linear map from that activity to each voxel’s response.

Then test the map on stimuli it hasn’t seen. If the network’s activity predicts the brain’s response to new sentences, the two must share some of the information they carry.

How a brain is compared with a model. The same sentence goes to a language model and to a person in an fMRI scanner. The model's internal activity, a vector of numbers, is mapped to each voxel's response by a linear regression fitted on training sentences. On new sentences the predicted response is compared with the measured one, and the correlation is divided by the noise ceiling: how well the brain predicts itself across repeated measurements or across people How a brain is compared with a model. The same sentence goes to a language model and to a person in an fMRI scanner. The model's internal activity, a vector of numbers, is mapped to each voxel's response by a linear regression fitted on training sentences. On new sentences the predicted response is compared with the measured one, and the correlation is divided by the noise ceiling: how well the brain predicts itself across repeated measurements or across people

The prediction is scored as a correlation, and since the brain itself is noisy, the correlation is compared with a noise ceiling: how well one recording of the same person, or other people, predicts another.

“The model explains 80% of the explainable variance” means its correlation is 80% of that ceiling. The other common method, representational similarity analysis, skips the map: it compares which stimuli the brain treats as similar with which ones the network treats as similar[2].

Both methods answer a narrower question than it seems. They show that two systems carry overlapping information about the stimuli, not that they compute it the same way.

A linear map is flexible enough to find shared information in very different machinery. And the details of the test matter enormously, as we will see.

Three numbers to keep in mind when reading these studies

Raw correlations are small. In one of the best-known language studies the best models reached a correlation of about 0.05 with single fMRI voxels, against a ceiling of 0.06[3]. “Close to the ceiling” can mean that both numbers are close to zero.

The ceiling is sometimes extrapolated. “Nearly 100% of explainable variance” in Schrimpf et al.’s much-cited paper is a ratio of correlations, with ceilings of r = 0.17 to 0.32, estimated for an infinite number of participants[4].

How the data are split decides a lot. If sentences from the same passage land in both the training and the test set, a model can score well by recognising the passage rather than the sentence.

A 2026 re-analysis found that several of the field’s best-known results depend on such splits, and that a model knowing only a sentence’s position in the experiment matches GPT-2 on them[5].

Language models and the brain’s language network

The human brain has a set of regions that respond when we read or hear language and barely respond to anything else: not to arithmetic, logic, music or code.

Evelina Fedorenko’s group calls it the language network; it sits mostly in the left hemisphere, along the temporal lobe and in the inferior frontal gyrus. Here it is as an atlas of 804 people: the colour shows how often each spot of cortex was among a person’s most language-selective.

The language network on an inflated cortex, both hemispheres seen from the side. Violet marks where, across 804 people, a spot was among the most responsive to sentences compared with lists of nonwords. The left hemisphere has large, strong patches in the inferior and middle frontal gyri and along the temporal lobe from front to back; the right hemisphere has weaker, smaller versions of the same patches The language network on an inflated cortex, both hemispheres seen from the side. Violet marks where, across 804 people, a spot was among the most responsive to sentences compared with lists of nonwords. The left hemisphere has large, strong patches in the inferior and middle frontal gyri and along the temporal lobe from front to back; the right hemisphere has weaker, smaller versions of the same patches

Since 2021 one result has been replicated with fMRI, magnetoencephalography and electrodes placed on the brains of epilepsy patients: the internal activity of language models predicts this network’s response to language better than any other model we have.

Schrimpf and colleagues compared 43 models and found transformers best; across models, the ones better at predicting the next word predicted the brain better, while performance on language-understanding benchmarks didn’t[4].

Caucheteux and King trained thousands of transformers from scratch and found brain similarity rising with the ability to predict words from context, with a twist: it peaked before the models were at their best, and “the highest brain scores are not achieved by the very best language transformers”[3].

The fit also scales. Training encoding models on about twenty hours of fMRI per person, Antonello, Vaidya and Huth found that predictions improved roughly logarithmically with model size, by about 15% from 125 million to 30 billion parameters, and then levelled off[6].

In some regions the best models came close to the ceiling. And it doesn’t take an internet of text: GPT-2 trained on 100 million words, about what a child hears in ten years, predicted the brain as well as GPT-2 trained on billions[7].

Here the honest part starts. What is the network actually matching? Kauf and colleagues scrambled the word order of sentences and lost only 8 to 24% of the predictivity; removing the content words lost up to 83%[8].

What the models share with the language network is mostly the meaning of words, plus some of how words combine, and much less syntax than the word “language model” suggests.

The 2026 re-analysis went further: on the classic datasets, a sentence’s position, the rate of words and static word embeddings together accounted for over 80% of GPT-2-XL’s predictivity[5].

The most attractive story in the field, that the brain is a next-word predictor because next-word predictors fit it best, has been pulled from both sides.

Goldstein and colleagues recorded from 1,339 electrodes while nine patients listened to a podcast and found information about the upcoming word up to 800 milliseconds before it was heard, and surprise after it, much as a language model would show[9].

In 2026 Schönmann and colleagues showed that the same signature appears when the analysis is run on the speech sound itself, or on word embeddings that cannot predict anything, because words in natural speech are statistically dependent on their neighbours[10].

Goldstein’s team had warned about this themselves: “These shared computational principles, however, do not imply that the human brain and DLMs implement these computations in a similar way.”

Language models are the best encoding models of the language network we have, and what they mostly capture is word meaning. Whether the brain computes language by predicting the next word is still open.

Language and thought rely on different networks

The language network has a second property that matters more than any correlation. People who lose it, after a stroke that leaves them with global aphasia, unable to understand or produce sentences, can still do arithmetic, solve logic problems and reason about what other people believe[11].

In healthy people the network stays quiet during arithmetic, logic and music. In the brain, language and thought are different systems that talk to each other[12].

Mahowald, Ivanova and colleagues turned this into a distinction for language models: formal competence, knowing the rules and patterns of a language, and functional competence, using language to reason, plan and know about the world.

A fluent text tempts us to infer the second from the first, the “good at language, therefore good at thought” fallacy[13]. The brain data agree.

In a family of models trained from scratch, similarity to the language network peaked after only 2 to 8 billion tokens of training, while reasoning and knowledge kept improving through 300 billion[14].

So a large language model is two things at once. Part of it is a good model of the human language network. The rest, the arithmetic, the facts, the code, has to live somewhere else in the brain, and the language network doesn’t predict which.

Decoding language from brain scans

The same models made it possible to decode language from brain activity. Tang, LeBel, Jain and Huth decoded the gist of stories that people heard or imagined, and of silent films they watched, from fMRI alone, using a GPT model to propose candidate word sequences and an encoding model to pick the one that best explained the scan[15].

The output is paraphrase, not transcript: most words come out wrong, and the meaning comes through. The decoder had to be trained on about 16 hours of each person’s own scans, didn’t work across people, and failed when the person resisted by thinking of something else.

The most practical use is in medicine. Electrodes implanted in the speech-motor cortex of a man with ALS let him speak through a computer at about 32 words a minute with 97.5% accuracy, over 248 hours of everyday use[16].

A language model corrects the raw output into likely sentences. The signal is the attempt to move the mouth, not the language network and not thought.

Vision models and the visual cortex

The visual system was where this whole line of work began, and where it is most mature.

Visual information flows from the primary visual cortex at the back of the head forward along the underside of the brain, through areas that respond to edges, then to shapes and textures, then to whole objects and faces in the inferior temporal cortex. This ventral stream is a hierarchy, like the stack of layers in a convolutional network.

The left hemisphere seen from the side, from the middle and from below, with the ventral visual stream in three shades: early visual cortex at the back of the head and along the calcarine sulcus, middle stages on the lateral occipital surface, and late stages in the fusiform and inferior temporal gyri on the underside of the temporal lobe The left hemisphere seen from the side, from the middle and from below, with the ventral visual stream in three shades: early visual cortex at the back of the head and along the calcarine sulcus, middle stages on the lateral occipital surface, and late stages in the fusiform and inferior temporal gyri on the underside of the temporal lobe

In 2014 Yamins, DiCarlo and colleagues showed that the resemblance is more than a picture. They trained networks only to recognise objects, never fitting them to neural data, and found that the better a network categorised, the better its top layer predicted single sites in the inferior temporal cortex of monkeys, and an intermediate layer predicted area V4.

The best model predicted 48.5% of the explainable variance of IT, nearly twice as much as any model built before it[17]. A system optimised for a task had converged, in part, on the representation that evolution and development produced for the same task.

What followed is the most instructive turn in the field. If accuracy made networks brain-like, more accurate networks should be more brain-like still.

For a while they were. Then four groups with different data found the payoff gone.

Five studies in order. Yamins and colleagues in 2014 found that a network trained only to categorise objects predicted 48.5% of the explainable variance of monkey inferior temporal cortex. Brain-Score in 2018 found the link between ImageNet accuracy and brain similarity weakening above about 70% accuracy. Linsley and colleagues in 2023 found that across 135 networks more accurate ones were worse models of IT. Gokce and Schrimpf in 2025 found, over more than 600 models, that agreement with behaviour kept improving with scale while agreement with neurons saturated. Conwell and colleagues in 2024 found ImageNet accuracy unrelated to how well 224 models predicted human fMRI Five studies in order. Yamins and colleagues in 2014 found that a network trained only to categorise objects predicted 48.5% of the explainable variance of monkey inferior temporal cortex. Brain-Score in 2018 found the link between ImageNet accuracy and brain similarity weakening above about 70% accuracy. Linsley and colleagues in 2023 found that across 135 networks more accurate ones were worse models of IT. Gokce and Schrimpf in 2025 found, over more than 600 models, that agreement with behaviour kept improving with scale while agreement with neurons saturated. Conwell and colleagues in 2024 found ImageNet accuracy unrelated to how well 224 models predicted human fMRI

Brain-Score, a benchmark that scores networks against monkey recordings, found the link between ImageNet accuracy and brain similarity weakening above about 70% accuracy[18]. Across 135 networks, Linsley and colleagues found that the more accurate ones had become worse models of IT[19].

Over 600 models trained under controlled conditions, alignment with behaviour kept improving with scale while alignment with neurons saturated[20]. And in human fMRI from the Natural Scenes Dataset, ImageNet accuracy didn’t predict how brain-like a model was at all: the rank correlation was −0.006 by one measure and 0.17 by another[21].

Conwell’s study, 224 models and 1.8 billion regressions, found something even more uncomfortable.

Convolutional networks and transformers trained the same way scored r = 0.67 and 0.66; 126 of the 224 models scored within 0.1 of each other; models whose representations barely correlated with each other got the same brain score; and the same models explained about 31% or about 79.5% of the explainable variance, depending on which of two reasonable versions of the analysis was used.

What mattered most was what images a model was trained on. The authors’ conclusion is about our instruments as much as about brains: “standard methods used to link models to brains may be too flexible.”

Task-trained networks predicted the visual cortex better than any hand-built model. Since then, accuracy has stopped buying similarity, architecture barely matters, and the measuring stick is looser than the scores suggest.

How human and machine vision differ

The differences show up in behaviour. On images with the shape of one object and the texture of another, a cat’s outline filled with elephant skin, people name the shape 95.9% of the time; ImageNet-trained networks mostly name the texture[22].

In an odd-one-out task that needs the 3D shape of an object seen from different angles, people score 78% and the best vision model 44%[23]. Frontier multimodal models of 2024 solved a classic mental rotation test 33% of the time, against 78.5% for people and 25% by chance[24].

One difference turned out smaller than it looked. Tiny, invisible changes to an image can make a network see a different object, which seemed to prove that networks are fragile and brains are not.

But when the responses of monkey IT neurons were attacked the same way, they turned out about as fragile as networks trained to resist such attacks[25].

People resist because they get to look again; a 63-millisecond glimpse followed by a mask can be nudged by the same perturbations that fool networks[26].

Reconstructing images from brain activity

If you have seen images “reconstructed from brain scans”, they came from the Natural Scenes Dataset: eight people, each lying in a 7-tesla scanner for 30 to 40 sessions over a year and looking at 9,000 to 10,000 photographs[27].

Pipelines like MindEye2 map the scans into the input space of an image generator. Asked to pick which of 300 photographs a person saw, MindEye2 is right 98.3% of the time; with only one hour of a new person’s data, on top of 40 hours each from seven others, 57.4%[28].

Retrieval is real. Reconstruction is mostly the generator’s. On a dataset built so that test images come from categories absent from training, the same methods produced realistic pictures of the training categories instead of what people saw; geometric shapes came out as photographs.

The authors call what these methods do “a combination of ‘classification’ and ‘hallucination’“[29]. The brain supplies a rough steer, and a model trained on billions of images supplies the picture.

Place cells, grid cells and the hippocampus

The hippocampus and the entorhinal cortex next to it hold the brain’s best-understood map. Place cells fire when an animal is in one particular spot; grid cells fire at every vertex of a triangular lattice that tiles the whole environment. John O’Keefe and May-Britt and Edvard Moser shared the 2014 Nobel Prize for finding them[30].

A coronal slice of the brain with both hippocampi highlighted, and two idealised firing maps over a one-metre box: a place cell fires in one spot, a grid cell at every vertex of a triangular lattice A coronal slice of the brain with both hippocampi highlighted, and two idealised firing maps over a one-metre box: a place cell fires in one spot, a grid cell at every vertex of a triangular lattice

In 2018 a DeepMind team trained a recurrent network to keep track of its position from its own movements, and grid-like units appeared in it, about a quarter of the units in one layer, with scales like a rat’s; an agent using them navigated mazes and took shortcuts[31].

It looked like a task explaining a cell type. Then Schaeffer, Khona and Fiete trained more than 11,000 such networks.

Most learned to keep track of position; fewer than 10% showed grid-like units, only with one particular, hand-chosen way of reading out the position, and never with the several grid modules of different scales that real brains have[32].

Banino’s networks had needed dropout too. The grid came partly from the task and partly from the modeller’s choices.

A model can also look as if it has a map without having one. Transformers trained on turn-by-turn taxi rides in Manhattan chose a valid next turn almost every time and often found shortest paths, yet the street map implied by their behaviour was impossible, with streets running in orientations the city doesn’t have.

Replace 1% of their turns with random detours, and only 69% of their routes stayed valid; replace 10%, and 8% did[33].

Memory systems and the context window

The brain has several memory systems, and a language model has several stores, and it is tempting to pair them. The pairing is a useful map if you remember that it is a map.

Human memory systems paired with a language model's stores. Working memory, about four chunks kept active, pairs with the context window, a hundred thousand tokens or more held passively and used worst in the middle. Episodic memory, one-shot memories of one's own experiences, pairs with retrieval of documents, which are not the model's own past. Semantic memory, slowly consolidated facts, pairs with the weights, absorbed slowly in training. Consolidation by replay during sleep pairs with nothing: what happens in a conversation never reaches the weights Human memory systems paired with a language model's stores. Working memory, about four chunks kept active, pairs with the context window, a hundred thousand tokens or more held passively and used worst in the middle. Episodic memory, one-shot memories of one's own experiences, pairs with retrieval of documents, which are not the model's own past. Semantic memory, slowly consolidated facts, pairs with the weights, absorbed slowly in training. Consolidation by replay during sleep pairs with nothing: what happens in a conversation never reaches the weights

The weights of a model hold what it absorbed slowly in training, distributed and not individually addressable, which is loosely like semantic memory: knowing that Paris is in France without remembering when you learned it.

The context window holds the current conversation, available all at once and lost when it ends, which is loosely like working memory. Retrieval systems that search a database and paste documents into the context are loosely like recalling an episode.

Each analogy breaks somewhere instructive. Human working memory holds about four chunks of meaning at a time[34]; a context window holds a hundred thousand tokens and more.

But models use long contexts unevenly: information in the middle is used worse than at the beginning or the end, a U-shaped curve that resembles the human serial-position effect[35].

Retrieval brings back documents, not the system’s own experiences. And attention itself is, mathematically, one update of an associative memory, the kind of network neuroscientists have long used as a model of memory recall[36]; that is a statement about equations, not about the hippocampus.

The deepest difference is consolidation. In the brain, what happens during the day is replayed during sleep: cells that fired together during a task fire together again, in sequences compressed about twentyfold (Lee & Wilson, 2002).

Suppressing the ripples that carry this replay impairs memory, and prolonging them improves it[37, 38]. A deployed language model has no such step: nothing it learns in a conversation is written back into its weights.

Training it further on new facts erases old ones, and standard deep networks trained continually slowly lose the ability to learn at all[39]. Brains keep learning for a lifetime.

Dopamine and reinforcement learning

The cleanest case of a match runs in an unexpected direction. In the 1980s, reinforcement-learning researchers developed temporal-difference learning, in which an agent learns to predict future reward from the error between what it predicted and what it got.

In 1997 Schultz, Dayan and Montague showed that dopamine neurons in the monkey midbrain fire exactly like that error[40]. An unexpected drop of juice makes them burst.

Once a light reliably predicts the juice, the burst moves to the light and the juice itself gets no response. If the juice then fails to arrive, the neurons fall silent at the moment it was due.

Three panels of the prediction error of a temporal-difference model across one trial. Before learning, the error is a positive spike when the juice arrives. After learning, the spike has moved to the cue that predicts the juice, and there is nothing at the juice itself. When the juice is withheld after learning, there is a spike at the cue and a negative dip at the moment the juice was expected. These are the three patterns Schultz and colleagues recorded in dopamine neurons of monkeys Three panels of the prediction error of a temporal-difference model across one trial. Before learning, the error is a positive spike when the juice arrives. After learning, the spike has moved to the cue that predicts the juice, and there is nothing at the juice itself. When the juice is withheld after learning, there is a spike at the cue and a negative dip at the moment the juice was expected. These are the three patterns Schultz and colleagues recorded in dopamine neurons of monkeys

The chart above is not a recording. It is a temporal-difference model computed for this post, and it reproduces all three patterns that the monkeys’ neurons showed.

The story went further. In AI, distributional reinforcement learning predicts the whole range of possible rewards rather than their average, with some predictors optimistic and some pessimistic.

Dabney and colleagues predicted that dopamine neurons would show the same diversity, and found it in recordings from mice: different neurons flipped from silence to firing at different reward sizes[41].

In 2025 two more papers found dopamine neurons working at many timescales at once and encoding a map of both when and how much reward to expect[42, 43].

Two caveats keep this honest. The distributional result rests on tens of cells in five mice. And the prediction-error account itself has a serious challenger: a careful 2022 study concluded that dopamine release in the nucleus accumbens tracks causal associations rather than prediction errors[44].

The analogy that people draw most often, between dopamine and the reinforcement learning that turns a language model into a chat assistant, is the weakest. Both use a single number to say “more of this”.

But in RLHF that number comes from a second model trained to guess human preferences, and optimising against it too hard makes answers worse[45] and more sycophantic[46]. It also damaged GPT-4’s sense of its own reliability: before this training its confidence matched its accuracy almost perfectly, after it, ten times worse[47].

Intuitive physics, robots and Moravec’s paradox

In 1988 the roboticist Hans Moravec noticed that what is hard for people, chess, logic, exams, was easy for computers, and what a one-year-old does without thinking, seeing and moving, was the hardest of all. Nearly forty years of progress have kept the paradox and given it numbers.

Video models trained to predict the hidden parts of natural videos learn some intuitive physics: a model called V-JEPA tells a possible video from an impossible one, a ball that vanishes behind a screen or jumps from one place to another, 98% of the time on a classic test borrowed from infant research[48].

On a harder version of the same test with more complex scenes, people are right 96.4% of the time and the best models are close to a coin flip.

Accuracy at telling possible from impossible videos on the IntPhys 2 benchmark. People, by majority of three annotators, 96.4%. The best models are near chance at 50%: V-JEPA 2 57.5%, Gemini 2.5 Flash 55.6%, GPT-4o 53.8%, V-JEPA with RoPE 53.8%, Qwen-VL 2.5 52.3%, Cosmos-4B 49.4% Accuracy at telling possible from impossible videos on the IntPhys 2 benchmark. People, by majority of three annotators, 96.4%. The best models are near chance at 50%: V-JEPA 2 57.5%, Gemini 2.5 Flash 55.6%, GPT-4o 53.8%, V-JEPA with RoPE 53.8%, Qwen-VL 2.5 52.3%, Cosmos-4B 49.4%

Robots now borrow knowledge from the web. RT-2, a vision-language model fine-tuned to output robot actions, doubled the success of its predecessor on objects and settings it hadn’t seen, from 32% to 62%, and could follow instructions like “pick up the object that could be used as an improvised hammer”.

Its authors were clear about the limit: “the robot does not acquire any ability to perform new motions”[49].

V-JEPA 2, used as a world model to plan a robot arm’s movements, grasped a cup 65% of the time and a box 25%, spending 16 seconds planning each move[50]. No benchmark measures human and robot dexterity on the same tasks, which says something by itself.

The brain-model link holds in motor control too, for a body simulated in detail. A network trained to drive a realistic virtual rat, imitating the movements of real rats, produced activity that predicted the neurons of real rats’ motor cortex and striatum better than any measured feature of their movements did[51].

Which brain systems have an AI counterpart

Put the evidence together and a map appears. Colour each brain system by how close the best AI counterpart is, judged by the studies above: strong where models predict the system’s activity or reproduce its signature, partial where they do some of its job without its properties, and absent where there is nothing to compare.

The left hemisphere from the side and from the middle, coloured by how close an AI counterpart is. Visual cortex, auditory cortex and the language network: a close counterpart. Motor cortex, dorsolateral prefrontal cortex and the parahippocampal region: a partial one. The insula, the anterior cingulate and the ventromedial prefrontal cortex, which handle the body's state and emotion: none yet The left hemisphere from the side and from the middle, coloured by how close an AI counterpart is. Visual cortex, auditory cortex and the language network: a close counterpart. Motor cortex, dorsolateral prefrontal cortex and the parahippocampal region: a partial one. The insula, the anterior cingulate and the ventromedial prefrontal cortex, which handle the body's state and emotion: none yet

The table names the systems behind the colours and the counterpart each was compared with.

A table of brain systems, their nearest AI counterparts and the state of the match. Strong: vision and vision networks, which predict the ventral stream though gains have stopped; hearing and speech and models like Whisper, which predict the auditory-to-language hierarchy; the language network and language models, which predict it mostly through word meaning; the dopamine teaching signal and temporal-difference learning, which predicted new properties of neurons. Partial: the hippocampus and space, where grid-like units appear only by design and there is no consolidation; motor cortex and robot policies, which gained web knowledge but no new motions; prefrontal control and reasoning models, without a brain-level match. Absent: the insula and interoception, with no body to sense, and sleep and neuromodulators, with no counterpart A table of brain systems, their nearest AI counterparts and the state of the match. Strong: vision and vision networks, which predict the ventral stream though gains have stopped; hearing and speech and models like Whisper, which predict the auditory-to-language hierarchy; the language network and language models, which predict it mostly through word meaning; the dopamine teaching signal and temporal-difference learning, which predicted new properties of neurons. Partial: the hippocampus and space, where grid-like units appear only by design and there is no consolidation; motor cortex and robot policies, which gained web knowledge but no new motions; prefrontal control and reasoning models, without a brain-level match. Absent: the insula and interoception, with no body to sense, and sleep and neuromodulators, with no counterpart

The strongest matches are the sensory and language systems, and the dopamine teaching signal. They have something in common: each was matched by a model trained on a large amount of the same kind of input the brain gets, to do a similar job, and the similarity grew out of the training rather than being built in.

The weakest are the systems that run the body, keep track of the self and keep learning over a lifetime.

Is the brain more efficient than AI?

This is the question where the most nonsense is written, so the numbers here come with their units and their uncertainty.

Energy

The usual figure is that a brain runs on 20 watts. That is the energy of the glucose it burns; positron emission tomography puts it nearer 17 W, and only about 6 W of that becomes ATP that can do work.

Most of the brain’s energy goes into moving signals rather than computing with them: in an audit of the human cortex, long-distance communication cost about 3.5 W and computation in the narrow sense, at the synapses, about 0.1 W[52]. Computers have the same problem, where moving data costs more energy than the arithmetic on it.

Training is where AI looks extravagant. Twenty years of a 20-watt brain come to about 3.5 megawatt-hours. Training GPT-3 took 1,287 MWh[53], Llama 3.1 405B 30.84 million GPU-hours, which comes to about 22 GWh for the GPUs alone[54], and Epoch AI estimates Grok 4 at about 310 GWh[55].

Two bar charts on log scales. To train: a brain over 20 years at 20 watts uses 3.5 megawatt-hours, GPT-3's training 1,287, Llama 3.1 405B about 21,600 for its GPUs alone, Grok 4 an estimated 310,000. To answer: a brain in one minute uses 0.33 watt-hours, a median Gemini prompt 0.24, a typical GPT-4o query 0.3, a GPT-4o query with 10,000 tokens of input 2.5 and with 100,000 tokens about 40 Two bar charts on log scales. To train: a brain over 20 years at 20 watts uses 3.5 megawatt-hours, GPT-3's training 1,287, Llama 3.1 405B about 21,600 for its GPUs alone, Grok 4 an estimated 310,000. To answer: a brain in one minute uses 0.33 watt-hours, a median Gemini prompt 0.24, a typical GPT-4o query 0.3, a GPT-4o query with 10,000 tokens of input 2.5 and with 100,000 tokens about 40

But training is paid once and shared by everyone who uses the model afterwards. An answer from a chatbot is cheap: Google measured 0.24 watt-hours for a median Gemini text prompt in 2025, after a 33-fold drop in a year[56], and Epoch AI estimates about 0.3 Wh for a typical GPT-4o query[57].

That is the energy a brain uses in about 45 seconds. Long prompts and reasoning cost much more.

Per operation the comparison is speculative, because a brain doesn’t do floating-point operations. The most careful estimate puts the computation needed to match a brain at about 10¹⁵ operations per second, with a plausible range from 10¹³ to 10¹⁷[58].

On the central estimate, a brain does about 5 × 10¹³ operations per joule and an H100 GPU at its peak about 1.4 × 10¹², so the brain is some tens of times more efficient; across the range, anywhere from three times worse to a few thousand times better.

Data

Here the gap is larger and much less uncertain. A child hears somewhere between ten and a hundred million words before adolescence[59]; Llama 3.1 was trained on 15.6 trillion tokens, each a little less than a word[54]. That is five to six orders of magnitude.

Words and tokens on a log scale. A child hears between 6 and 60 million words by age 5, up to about 100 million by age 13, and 30 to 400 million words by age 20 including reading. The BabyLM challenge allows 100 million words. GPT-3 was trained on 500 billion tokens and Llama 3.1 on 15.6 trillion Words and tokens on a log scale. A child hears between 6 and 60 million words by age 5, up to about 100 million by age 13, and 30 to 400 million words by age 20 including reading. The BabyLM challenge allows 100 million words. GPT-3 was trained on 500 billion tokens and Llama 3.1 on 15.6 trillion

The comparison is not quite fair to the child, who also sees, touches, moves and is taught by people who care, and arrives with priors that evolution spent a long time on. It is also not quite fair to the model, which has to learn from text alone what a child gets from the world.

The BabyLM challenge trains models on a child’s budget of 100 million words: the best come within a few points of people on tests of grammar and stay near chance on world knowledge[60].

Speed and breadth

The models win where machines always win. A neuron fires at most a few hundred times a second; a model writes a page while a person reads a sentence. A model can be copied a million times, read a library in a day and remember the details.

One synapse stores about 4.7 bits[61], one parameter of a language model about 3.6 bits[62], and a large model has hundreds to thousands of times fewer parameters than a brain has synapses while knowing far more facts than any person.

Per answer, a chatbot costs about a minute of brain time. Per lesson learned, the brain is cheaper by orders of magnitude in data, and probably in energy too.

What AI models still lack

A list of what the brain has and AI models don’t would be long. Five items on it come up again and again in the evidence.

Learning that never stops. People learn from one example and keep learning for decades without forgetting what they knew. Models learn from millions of examples, in a separate training phase, and forget when trained further[39].

The brain’s solution seems to involve two learning systems, a fast hippocampus and a slow cortex, and sleep to move memories between them[63]; nothing like that runs in a deployed model.

A body. The systems that the models match best, language and vision, are the ones that can be trained from data on the internet. The ones they match worst, movement, intuitive physics and space, are the ones a body learns by acting.

Wanting something. A brain is driven by needs, curiosity and a reward system tuned by evolution. A model optimises what its trainers chose, and does nothing when nobody asks it anything.

Knowing what it doesn’t know. People are imperfect judges of their own knowledge, but they can usually tell a guess from a memory. Models trained to please are confidently wrong, and users trust long, fluent explanations more than they should[64].

Sleep and neuromodulation. The brain changes its own learning rate, attention and mood with a handful of chemicals, and reorganises its memories every night. There is no clear counterpart in any current AI system.

Whether there is something it is like to be a model, the question of consciousness, is a separate story, and a hard one.

The most careful attempt so far, which checked current systems against the signatures that the leading scientific theories of consciousness predict, concluded that no current AI system is conscious and that there is no obvious technical barrier to building one that would meet those signatures[65].

What the comparison shows

Five conclusions seem safe.

First, the convergence is real. Networks trained only to do what the brain does, recognise objects, predict words, predict reward, have repeatedly turned out to resemble the brain more than models designed to resemble it.

The likeliest reason is not that engineers copied the brain, but that the same problem, solved from the same kind of data, pushes any good solution toward similar representations.

Second, the resemblance is mostly in what is represented, not how it is computed. A transformer reads a whole sentence at once; the brain hears it one word after another. Linear maps find shared information in very different machinery.

Third, the similarity scores have saturated, and our methods are part of the reason. Better models stopped becoming more brain-like in vision, and in language much of the measured similarity is word meaning and statistics of the experiment.

Fourth, a language model is a model of one part of the brain. It matches the language network, which in people is not where thinking happens, and everything else it does has to be compared with other systems, one at a time.

Fifth, the brain is not a slower GPU. It learns from far less, on less energy, all its life, inside a body, and the systems where models are furthest from it are the ones that make that possible.

The figures, the TD-learning model behind the dopamine chart and the preparation of the language atlas as code: articles/brain_ai.

References

  1. Tuckute et al., 2024. doi.org. ↩
  2. Khaligh-Razavi & Kriegeskorte, 2014. doi.org. ↩
  3. Caucheteux & King, 2022. doi.org. ↩
  4. Schrimpf et al., 2021. doi.org. ↩
  5. Hadidi et al., 2026. doi.org. ↩
  6. Antonello et al., 2023. arxiv.org. ↩
  7. Hosseini et al., 2024. doi.org. ↩
  8. Kauf et al., 2024. doi.org. ↩
  9. Goldstein et al., 2022. doi.org. ↩
  10. Schönmann et al., 2026. doi.org. ↩
  11. Fedorenko & Varley, 2016. doi.org. ↩
  12. Fedorenko et al., 2024. doi.org. ↩
  13. Mahowald et al., 2024. doi.org. ↩
  14. AlKhamissi et al., 2025. arxiv.org. ↩
  15. Tang et al., 2023. doi.org. ↩
  16. Card et al., 2024. doi.org. ↩
  17. Yamins et al., 2014. doi.org. ↩
  18. Schrimpf et al., 2018. doi.org. ↩
  19. Linsley et al., 2023. arxiv.org. ↩
  20. Gokce & Schrimpf, 2025. arxiv.org. ↩
  21. Conwell et al., 2024. doi.org. ↩
  22. Geirhos et al., 2019. arxiv.org. ↩
  23. Bonnen et al., 2024. arxiv.org. ↩
  24. Ramakrishnan et al., 2025. arxiv.org. ↩
  25. Guo et al., 2022. arxiv.org. ↩
  26. Elsayed et al., 2018. arxiv.org. ↩
  27. Allen et al., 2022. doi.org. ↩
  28. Scotti et al., 2024. arxiv.org. ↩
  29. Shirakawa et al., 2025. arxiv.org. ↩
  30. Hafting et al., 2005. doi.org. ↩
  31. Banino et al., 2018. doi.org. ↩
  32. Schaeffer et al., 2022. doi.org. ↩
  33. Vafa et al., 2024. arxiv.org. ↩
  34. Cowan, 2001. doi.org. ↩
  35. Liu et al., 2024. doi.org. ↩
  36. Ramsauer et al., 2021. arxiv.org. ↩
  37. Girardeau et al., 2009. doi.org. ↩
  38. Fernández-Ruiz et al., 2019. doi.org. ↩
  39. Dohare et al., 2024. doi.org. ↩
  40. Schultz et al., 1997. doi.org. ↩
  41. Dabney et al., 2020. doi.org. ↩
  42. Masset et al., 2025. doi.org. ↩
  43. Sousa et al., 2025. doi.org. ↩
  44. Jeong et al., 2022. doi.org. ↩
  45. Gao et al., 2023. arxiv.org. ↩
  46. Sharma et al., 2023. arxiv.org. ↩
  47. OpenAI, 2023. arxiv.org. ↩
  48. Garrido et al., 2025. arxiv.org. ↩
  49. Brohan et al., 2023. arxiv.org. ↩
  50. Assran et al., 2025. arxiv.org. ↩
  51. Aldarondo et al., 2024. doi.org. ↩
  52. Levy & Calvert, 2021. doi.org. ↩
  53. Patterson et al., 2021. arxiv.org. ↩
  54. Meta, 2024. arxiv.org. ↩
  55. Epoch AI, 2025. epoch.ai. ↩
  56. Elsworth et al., 2025. arxiv.org. ↩
  57. You, 2025. epoch.ai. ↩
  58. Carlsmith, 2020. openphilanthropy.org. ↩
  59. Frank, 2023. doi.org. ↩
  60. Hu et al., 2024. arxiv.org. ↩
  61. Bartol et al., 2015. doi.org. ↩
  62. Morris et al., 2025. arxiv.org. ↩
  63. Kumaran et al., 2016. doi.org. ↩
  64. Steyvers et al., 2025. doi.org. ↩
  65. Butlin et al., 2023. arxiv.org. ↩
]]> Eva Koroleva T5 and BART: the encoder–decoder, and a model that writes SQL 2026-05-15T14:00:00+00:00 2026-11-08T00:00:00+00:00 https://blog.korolevalab.com/blog/2026/transformer-architectures-05-t5 Transformer architectures: from BERT to diffusion · lesson 5 of 15 · all lessons · run on Kaggle · open in Colab · code on GitHub

In 1961 four researchers at MIT’s Lincoln Laboratory fed a computer questions on punched cards, such as “Where did the Red Sox play on July 7?” Their program, BASEBALL, turned each one into a list, Place = ? Team = Red Sox Month = July Day = 7, looked the answer up in the results of one American League season and printed it[1].

The authors had bigger plans: “the business executive, the military commander, and the scientist need to ask questions of the computer in ordinary English”.

Sixty-five years later, asking a database in English is a common feature, and the benchmark everyone measured it on, Spider, went from 12.4% of queries exactly right for the best model of 2018 to 86.6% of queries returning the right result for GPT-4 with careful prompting in 2023.

Then came Spider 2.0, built from real company databases “often containing over 1,000 columns”, and o1-preview solved 21.3% of its tasks [2].

A question in English and a query in SQL are two languages, and the architecture built for translating between languages is the encoder–decoder. Its most influential member is T5, Google’s model of 2019 that turned every task into text in, text out.

For a while encoder–decoders led Spider: in 2021 T5-3B with a constrained decoder reached 75.1% on the test set. Then decoder-only models took over. But T5 never left: its encoder turns your prompt into the conditioning of Stable Diffusion 3 and FLUX, and in 2025 Google turned Gemma into an encoder–decoder.

In this lesson you will write T5 yourself, load Google’s weights into it, look at what makes it different from BERT and GPT, and fine-tune it on Spider for 20 minutes on a free GPU. Then you will run every query it writes and see which ones are right.

A question and the database's schema go into T5's encoder, which reads all of it at once. The decoder writes an SQL query one token at a time, looking back at the encoder's output through cross-attention. The query is then run on the SQLite database, and its result is compared with the result of the correct query A question and the database's schema go into T5's encoder, which reads all of it at once. The decoder writes an SQL query one token at a time, looking back at the encoder's output through cross-attention. The query is then run on the SQLite database, and its result is compared with the result of the correct query

What you’ll do

  1. See how T5 turns every task into text in, text out, and make the smallest T5 translate, judge grammar and score similarity from a prefix.
  2. Take the encoder–decoder apart: an encoder that reads the whole input, a decoder that writes, and the cross-attention that connects them.
  3. Write T5 in PyTorch, with its unusual parts: a norm without a mean, attention without scaling and positions as 32 learned biases. Load Google’s weights into it and match Hugging Face’s outputs.
  4. Look at what the position buckets learned, and watch the decoder find the source words while it translates.
  5. Corrupt text the way T5 and BART were pre-trained, and let both models repair it.
  6. Teach T5 to write SQL for databases it has never seen, check every query by running it, and let SQLite pick the best of several attempts.

What you need: a GPU. On Kaggle, Settings → Accelerator → GPU; on Colab, Runtime → Change runtime type → T4 GPU. The notebook downloads t5-small and t5-base (1.1 GB), BART (560 MB) and the Spider databases (200 MB), and fine-tunes T5 for 20 minutes; on a T4 it runs in about 45 minutes.

Setup

One small helper package, tfs, draws the course’s plots and token chips and knows where the notebook runs (Kaggle, Colab or your own machine). Everything the lesson teaches is written out in the notebook itself.

import importlib.util
import subprocess
import sys

TFS = "git+https://github.com/xmarva/transformers-from-scratch@main"
if importlib.util.find_spec("tfs") is None:
    subprocess.check_call([sys.executable, "-m", "pip", "install", "-q", TFS])

from tfs import env
from tfs.viz import table, token_chips, use_course_style

env.quiet()  # no progress bars or environment warnings in the outputs
SEED = 42
env.seed_everything(SEED)
COLORS = use_course_style()
WORK_DIR = env.work_dir()

import json
import math
import random
import re
import sqlite3
import time
import zipfile
from collections import Counter

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import torch
import torch.nn as nn
import torch.nn.functional as F
from huggingface_hub import hf_hub_download
from matplotlib.colors import LinearSegmentedColormap
from safetensors.torch import load_file
from tfs.viz import heatmap
from transformers import AutoTokenizer
# the reference implementations, for the checks
from transformers import BartForConditionalGeneration, T5ForConditionalGeneration

SMOKE = env.smoke()  # CI runs the lesson on a small slice; you run it in full
device = "cuda" if torch.cuda.is_available() else "cpu"
assert device == "cuda" or SMOKE, "This lesson fine-tunes a model for 20 minutes: use a GPU"
if device == "cuda":
    print("GPU:", torch.cuda.get_device_name(0))

GPU: Tesla T4

1. Every task as text

In October 2019 a Google team published a 67-page study built on one idea: “Every task we consider—including translation, question answering, and classification—is cast as feeding our model text as input and training it to generate some target text”[3]. They named the model the Text-to-Text Transfer Transformer, T5.

A task is chosen by a prefix written in plain words, translate English to German: or cola sentence: for grammaticality, and a class is a word, acceptable or not acceptable. Even regression became text.

STS-B asks how similar two sentences are, on a scale up to 5; T5 rounded the scores to steps of 0.2 and wrote them out as strings like “3.8”, which recasts the regression as “a 21-class classification problem”.

The paper is remembered as much for its method as for the idea: on one setup it compared nearly every choice people argued about in 2019, from architectures and objectives to datasets, multi-task training and scale. Its dataset became a standard of its own.

C4, the Colossal Clean Crawled Corpus, is about 750 GB of English cut from one month of Common Crawl, April 2019, by blunt rules: keep only lines that end in a full stop or similar mark, drop pages with fewer than three sentences, pages with any word from a list of “Dirty, Naughty, Obscene or Otherwise Bad Words”, pages with “lorem ipsum”, and any page with a curly bracket, since “the curly bracket “{” appears in many programming languages […] but not in natural text”.

Code was not what C4 was for, and section 8 shows one thing that cost.

The original T5 checkpoints were pre-trained on a mixture of C4 and the paper’s supervised tasks, translation and summarisation among them. That is why the smallest of them still knows its prefixes. Here are four of them in five prompts, one model with no task-specific layer:

T5_SMALL = "google-t5/t5-small"
tokenizer = AutoTokenizer.from_pretrained(T5_SMALL)
reference = T5ForConditionalGeneration.from_pretrained(T5_SMALL).to(device).eval()

prompts = [
    "translate English to German: The guards arrive at the museum at seven in the morning.",
    "translate English to French: The night train to Vienna leaves at eleven.",
    "cola sentence: The guards arrives at the museum at seven.",
    "stsb sentence1: The night train leaves at eleven. "
    "sentence2: At eleven the night train departs.",
    "stsb sentence1: The night train leaves at eleven. "
    "sentence2: The museum opens at nine.",
]
batch = tokenizer(prompts, return_tensors="pt", padding=True).to(device)
with torch.no_grad():
    answers = reference.generate(**batch, max_new_tokens=40, num_beams=1, do_sample=False)
answers = tokenizer.batch_decode(answers, skip_special_tokens=True)
table(pd.DataFrame({"prompt": prompts, "answer": answers}))
promptanswer
translate English to German: The guards arrive at the museum at seven in the morning.Die Wachen kommen morgens um sieben Uhr im Museum an.
translate English to French: The night train to Vienna leaves at eleven.Le train de nuit à Vienne quitte à 11 heures.
cola sentence: The guards arrives at the museum at seven.acceptable
stsb sentence1: The night train leaves at eleven. sentence2: At eleven the night train departs.5.0
stsb sentence1: The night train leaves at eleven. sentence2: The museum opens at nine.0.0

The translations are right, and the similarity scores come back as text: 5.0 for a paraphrase and 0.0 for two sentences that have nothing in common. The grammar judgement is wrong.

“The guards arrives at the museum at seven” is called acceptable, a reminder that a prefix asks for a task and doesn’t promise the model can do it. This is the smallest T5, with 60 million parameters.

2. Two stacks and a bridge

BERT, in lesson 1, was an encoder: it read a whole text at once, every token looking at every other, and turned it into vectors. GPT, in lesson 4, was a decoder: it wrote text one token at a time, each token looking only at the ones before it.

The original Transformer of 2017 was both, because it was built for translation, where the input and the output are different texts. T5 and BART went back to that design.

The encoder–decoder. On the left the encoder: a stack of layers with self-attention that sees the whole input in both directions and a feed-forward network; it runs once and turns the input into a sequence of vectors. On the right the decoder: in every layer causal self-attention over what it has written so far, then cross-attention, whose queries come from the decoder and whose keys and values come from the encoder's output, then a feed-forward network. Only the decoder runs once per generated token The encoder–decoder. On the left the encoder: a stack of layers with self-attention that sees the whole input in both directions and a feed-forward network; it runs once and turns the input into a sequence of vectors. On the right the decoder: in every layer causal self-attention over what it has written so far, then cross-attention, whose queries come from the decoder and whose keys and values come from the encoder's output, then a feed-forward network. Only the decoder runs once per generated token

The encoder turns the \(S\) input tokens into a matrix \(H \in \mathbb{R}^{S \times d}\), once. The decoder has, in every layer, one part that neither BERT nor GPT had: cross-attention. It is ordinary attention with the queries taken from the decoder’s own stream \(Y \in \mathbb{R}^{T \times d}\) and the keys and values taken from the encoder’s output:

\[\begin{aligned} &\text{CrossAttention}(Y, H) \\ &\quad = \text{softmax}\!\left(\frac{(Y W_Q)(H W_K)^\top}{\sqrt{d_k}} \right) H W_V. \end{aligned}\]

The matrix of scores is \(T \times S\): each of the \(T\) positions being written looks at all \(S\) positions of the input, and no mask is needed, since the input is complete before the first word is written. Self-attention in the decoder stays causal, as in GPT.

Two consequences follow, and both matter for the rest of this lesson. The input is read with full context in both directions, as BERT reads, and it is read once: generating \(T\) tokens runs the encoder once and the decoder \(T\) times. And the parameters are split between two stacks: t5-small has 6 encoder layers and 6 decoder layers, so a token of output passes through only half of the model.

3. T5 in PyTorch

T5’s layer has the familiar shape, pre-norm with a residual around each part, as in GPT-2. Four details differ, and each is a decision the T5 paper made on purpose.

A norm without a mean. LayerNorm subtracts the mean of a vector, divides by the standard deviation, scales and shifts. T5’s norm only rescales:

\[\text{RMSNorm}(x) = g \odot \frac{x}{\sqrt{\frac{1}{d}\sum_{i=1}^{d} x_i^2 + \varepsilon}}.\]

Here \(g\) is a learned vector of scales and \(\varepsilon = 10^{-6}\) guards against division by zero. No mean is subtracted and no bias is added.

The idea has its own name and paper, RMSNorm [4], which argues that the rescaling, not the centring, is what makes LayerNorm work, and it is the norm of Llama, Qwen and nearly every model since. PyTorch has it built in, which makes for an easy check:

class Norm(nn.Module):
    def __init__(self, d, eps):
        super().__init__()
        self.weight = nn.Parameter(torch.ones(d))
        self.eps = eps

    def forward(self, x):
        variance = x.float().pow(2).mean(dim=-1, keepdim=True)
        x = x.float() * torch.rsqrt(variance + self.eps)
        return self.weight * x.to(self.weight.dtype)


x = torch.randn(2, 5, 8)
mine, builtin = Norm(8, 1e-6), nn.RMSNorm(8, eps=1e-6)
print("largest difference from nn.RMSNorm:", (mine(x) - builtin(x)).abs().max().item())

largest difference from nn.RMSNorm: 0.0

Attention without \(\sqrt{d_k}\). T5 drops the division by \(\sqrt{d_k}\) in the scores. Nothing is lost: Google’s initialisation draws the query matrix with a standard deviation \(\sqrt{d_k}\) times smaller than the key matrix’s, which has the same effect when training starts, and training adjusts it from there.

A transplant has to leave the division out too, or every score comes out eight times too small.

Positions as a bias. T5 has no position embeddings at all. Instead, each attention head adds a learned number to each score, which depends only on how far the key is from the query:

\[s_{ij} = q_i \cdot k_j + b_h\big(\text{bucket}(j - i)\big).\]

The distances are sorted into 32 buckets. In the encoder, which looks both ways, half of the buckets are for keys to the left and half for keys to the right.

Within each half, the first 8 distances get a bucket each, and the rest share the remaining 8 buckets on a logarithmic scale up to a distance of 128; anything farther falls into the last bucket.

Near words get precise positions, far ones only a rough sense of “far”. The decoder looks only back, so it uses all 32 buckets for distances to the left.

Written out for the encoder, with \(n = \lvert j - i \rvert\) the distance between key and query, the bucket within one half is

\[b(n) = \begin{cases} n, & n < 8, \\ \min\big(15,\; f(n)\big), & n \ge 8, \end{cases}\]

where

\[f(n) = 8 + \left\lfloor 8 \cdot \dfrac{\ln(n/8)}{\ln(128/8)} \right\rfloor,\]

and keys to the right of the query get \(b(n) + 16\). The logarithm spreads the eight far buckets so that each covers a range about 1.4 times wider than the one before: distances 8 to 11 share a bucket, then 12 to 15, 16 to 22, and the last one starts at 91 and takes everything farther.

The decoder uses the same rule with 16 exact distances and 16 far buckets, \(16 + \lfloor 16 \cdot \ln(n/16) / \ln(128/16) \rfloor\), capped at 31.

def position_bucket(relative, bidirectional, buckets=32, max_distance=128):
    bucket = torch.zeros_like(relative)
    if bidirectional:
        buckets = buckets // 2
        # keys to the right of the query get the upper half of the buckets
        bucket = bucket + (relative > 0).long() * buckets
        distance = relative.abs()
    else:
        distance = (-relative).clamp(min=0)  # the decoder only looks back
    exact = buckets // 2
    log_distance = torch.log(distance.clamp(min=1).float() / exact)
    scaled = log_distance / math.log(max_distance / exact)
    far = exact + (scaled * (buckets - exact)).long()
    far = far.clamp(max=buckets - 1)
    return bucket + torch.where(distance < exact, distance, far)


relative = torch.tensor([-300, -128, -64, -20, -9, -8, -7, -1,
                         0, 1, 7, 8, 20, 64, 128, 300])
table(pd.DataFrame({"key minus query": relative.tolist(),
                    "encoder bucket": position_bucket(relative, True).tolist(),
                    "decoder bucket": position_bucket(relative, False).tolist()}))
key minus queryencoder bucketdecoder bucket
-3001531
-1281531
-641426
-201017
-989
-888
-777
-111
000
1170
7230
8240
20260
64300
128310
300310

A distance of 1 and a distance of 7 get different buckets; 64 and 100 share one. Every key to the right of a decoder query lands in bucket 0 with the query itself, but the causal mask hides those keys anyway. The whole map, for distances up to 150 either way:

distances = torch.arange(-150, 151)
fig, ax = plt.subplots(figsize=(6.4, 2.8))
ax.step(distances, position_bucket(distances, True), where="mid", color=COLORS["enc"],
        label="encoder: both directions")
ax.step(distances, position_bucket(distances, False), where="mid", color=COLORS["dec"],
        label="decoder: only back")
ax.axvspan(-7.5, 7.5, color=COLORS["grid"], zorder=0)
ax.text(0, 33.5, "one bucket per distance", ha="center", color=COLORS["muted"], fontsize=9)
ax.set_xlabel("key position minus query position")
ax.set_ylabel("bucket")
ax.set_ylim(-1, 36)
ax.legend(loc="lower right", frameon=False)
plt.show()

Output of the code above

The table of biases, 32 buckets by the number of heads, lives in the first layer of the encoder and the first layer of the decoder. Every other layer reuses the same biases, so a model of six layers learns positions with \(2 \times 32 \times 8 = 512\) numbers. GPT-2 spent 786,432 on its position embeddings.

The output layer. T5’s decoder scores the vocabulary with the input embedding matrix, tied as in GPT-2, after multiplying the hidden state by \(d^{-1/2}\).

Google’s Mesh TensorFlow code does this whenever the input and output embeddings are shared, and Hugging Face’s code copies it. Leave it out and every logit is \(\sqrt{512} \approx 23\) times too large for t5-small.

The attention layer below takes an optional source: without it, it is self-attention; with it, the keys and values come from the encoder and it is cross-attention.

The optional cache is the KV cache of lesson 4, with one addition: in cross-attention the keys and values come from the encoder, which doesn’t change while the decoder writes, so they are computed once for the whole answer. For the pictures of section 5 it can also return the attention weights.

class Attention(nn.Module):
    def __init__(self, config, relative=False, bidirectional=True):
        super().__init__()
        d = config["d_model"]
        self.heads, self.head_dim = config["num_heads"], config["d_kv"]
        inner = self.heads * self.head_dim
        self.q = nn.Linear(d, inner, bias=False)
        self.k = nn.Linear(d, inner, bias=False)
        self.v = nn.Linear(d, inner, bias=False)
        self.o = nn.Linear(inner, d, bias=False)
        self.buckets = config["relative_attention_num_buckets"]
        self.bidirectional = bidirectional
        if relative:  # only the first layer has the table; the others reuse its biases
            self.relative_bias = nn.Embedding(self.buckets, self.heads)

    def position_bias(self, queries, keys):
        relative = keys[None, :] - queries[:, None]
        buckets = position_bucket(relative, self.bidirectional, self.buckets)
        return self.relative_bias(buckets).permute(2, 0, 1).unsqueeze(0)  # 1, heads, q, k

    def split(self, x):
        batch, length, _ = x.shape
        return x.view(batch, length, self.heads, self.head_dim).transpose(1, 2)

    def forward(self, x, source=None, bias=None, cache=None, keep_weights=False):
        q = self.split(self.q(x))
        if source is not None and cache is not None and "k" in cache:
            # cross-attention: the encoder's keys and values never change
            k, v = cache["k"], cache["v"]
        else:
            k = self.split(self.k(x if source is None else source))
            v = self.split(self.v(x if source is None else source))
            if source is None and cache is not None and "k" in cache:
                k = torch.cat([cache["k"], k], dim=2)
                v = torch.cat([cache["v"], v], dim=2)
        if cache is not None:
            cache["k"], cache["v"] = k, v
        weights = None
        if keep_weights:
            scores = q @ k.transpose(-2, -1) + bias  # no division by the square root of d_k
            weights = scores.float().softmax(dim=-1).to(q.dtype)
            y = weights @ v
        else:
            y = F.scaled_dot_product_attention(q, k, v, attn_mask=bias, scale=1.0)
        y = y.transpose(1, 2).reshape(x.shape[0], x.shape[1], -1)
        return self.o(y), weights

The feed-forward network of the original T5 is two matrices with a ReLU between them and no biases; there are no biases anywhere in T5. Dropout only matters for training, in section 8.

class FeedForward(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.wi = nn.Linear(config["d_model"], config["d_ff"], bias=False)
        self.wo = nn.Linear(config["d_ff"], config["d_model"], bias=False)
        self.drop = nn.Dropout(config["dropout_rate"])

    def forward(self, x):
        return self.wo(self.drop(F.relu(self.wi(x))))


class EncoderLayer(nn.Module):
    def __init__(self, config, first):
        super().__init__()
        d, eps = config["d_model"], config["layer_norm_epsilon"]
        self.norm1, self.attention = Norm(d, eps), Attention(config, relative=first)
        self.norm2, self.ff = Norm(d, eps), FeedForward(config)
        self.drop = nn.Dropout(config["dropout_rate"])

    def forward(self, x, bias):
        x = x + self.drop(self.attention(self.norm1(x), bias=bias)[0])
        return x + self.drop(self.ff(self.norm2(x)))


class DecoderLayer(nn.Module):
    def __init__(self, config, first):
        super().__init__()
        d, eps = config["d_model"], config["layer_norm_epsilon"]
        self.norm1 = Norm(d, eps)
        self.attention = Attention(config, relative=first, bidirectional=False)
        self.norm2, self.cross = Norm(d, eps), Attention(config)
        self.norm3, self.ff = Norm(d, eps), FeedForward(config)
        self.drop = nn.Dropout(config["dropout_rate"])

    def forward(self, y, memory, self_bias, cross_bias, cache=None, keep_weights=False):
        own = cache["self"] if cache is not None else None
        other = cache["cross"] if cache is not None else None
        y = y + self.drop(self.attention(self.norm1(y), bias=self_bias, cache=own)[0])
        out, weights = self.cross(self.norm2(y), source=memory, bias=cross_bias,
                                  cache=other, keep_weights=keep_weights)
        y = y + self.drop(out)
        return y + self.drop(self.ff(self.norm3(y))), weights

The model puts the stacks together. The encoder’s padding is hidden by a bias of \(-\infty\), in practice the smallest float, on the padded keys, both in the encoder and in cross-attention; the decoder adds the causal mask to its position biases.

decode takes past, the number of tokens already written, so that a single new token gets the right position biases when the earlier ones are in the cache.

def padding_bias(mask):
    blocked = torch.finfo(torch.float32).min
    return (1.0 - mask[:, None, None, :].float()) * blocked  # batch, 1, 1, keys


class T5(nn.Module):
    def __init__(self, config):
        super().__init__()
        d, eps = config["d_model"], config["layer_norm_epsilon"]
        self.d = d
        self.embed = nn.Embedding(config["vocab_size"], d)
        layers = config["num_layers"]
        self.encoder = nn.ModuleList(EncoderLayer(config, i == 0) for i in range(layers))
        self.encoder_norm = Norm(d, eps)
        layers = config.get("num_decoder_layers", config["num_layers"])
        self.decoder = nn.ModuleList(DecoderLayer(config, i == 0) for i in range(layers))
        self.decoder_norm = Norm(d, eps)
        self.drop = nn.Dropout(config["dropout_rate"])

    def encode(self, ids, mask):
        positions = torch.arange(ids.shape[1], device=ids.device)
        bias = self.encoder[0].attention.position_bias(positions, positions)
        bias = bias + padding_bias(mask)
        x = self.drop(self.embed(ids))
        for layer in self.encoder:
            x = layer(x, bias)
        return self.drop(self.encoder_norm(x))

    def decode(self, targets, memory, mask, caches=None, past=0, keep_weights=False):
        length = targets.shape[1]
        queries = torch.arange(past, past + length, device=targets.device)
        keys = torch.arange(past + length, device=targets.device)
        self_bias = self.decoder[0].attention.position_bias(queries, keys)
        future = keys[None, :] > queries[:, None]
        self_bias = self_bias.masked_fill(future, torch.finfo(torch.float32).min)
        cross_bias = padding_bias(mask)
        y = self.drop(self.embed(targets))
        cross = []
        for i, layer in enumerate(self.decoder):
            cache = caches[i] if caches is not None else None
            y, weights = layer(y, memory, self_bias, cross_bias, cache, keep_weights)
            cross.append(weights)
        y = self.drop(self.decoder_norm(y))
        logits = (y * self.d ** -0.5) @ self.embed.weight.T  # tied, and rescaled first
        return (logits, cross) if keep_weights else logits

4. Google’s weights in our code

The checkpoint’s names follow the structure of Google’s code, block.{i}.layer.{j}, where the parts of a layer are numbered: in the decoder, 0 is self-attention, 1 cross-attention and 2 the feed-forward network. We write down which goes where; load_state_dict refuses anything missing or left over.

def transplant(model, state):
    mine = {"embed.weight": state["shared.weight"],
            "encoder_norm.weight": state["encoder.final_layer_norm.weight"],
            "decoder_norm.weight": state["decoder.final_layer_norm.weight"]}
    for i in range(len(model.encoder)):
        old, new = f"encoder.block.{i}.layer.", f"encoder.{i}."
        for name in "qkvo":
            weight = state[old + f"0.SelfAttention.{name}.weight"]
            mine[new + f"attention.{name}.weight"] = weight
        mine[new + "norm1.weight"] = state[old + "0.layer_norm.weight"]
        mine[new + "ff.wi.weight"] = state[old + "1.DenseReluDense.wi.weight"]
        mine[new + "ff.wo.weight"] = state[old + "1.DenseReluDense.wo.weight"]
        mine[new + "norm2.weight"] = state[old + "1.layer_norm.weight"]
    for i in range(len(model.decoder)):
        old, new = f"decoder.block.{i}.layer.", f"decoder.{i}."
        for name in "qkvo":
            weight = state[old + f"0.SelfAttention.{name}.weight"]
            mine[new + f"attention.{name}.weight"] = weight
            weight = state[old + f"1.EncDecAttention.{name}.weight"]
            mine[new + f"cross.{name}.weight"] = weight
        for part in range(3):
            mine[new + f"norm{part + 1}.weight"] = state[old + f"{part}.layer_norm.weight"]
        mine[new + "ff.wi.weight"] = state[old + "2.DenseReluDense.wi.weight"]
        mine[new + "ff.wo.weight"] = state[old + "2.DenseReluDense.wo.weight"]
    for side in ["encoder", "decoder"]:
        key = f"{side}.block.0.layer.0.SelfAttention.relative_attention_bias.weight"
        mine[f"{side}.0.attention.relative_bias.weight"] = state[key]
    model.load_state_dict(mine)
    return model


def load_t5(repo):
    with open(hf_hub_download(repo, "config.json")) as f:
        config = json.load(f)
    state = load_file(hf_hub_download(repo, "model.safetensors"))
    return transplant(T5(config), state), config


t5, t5_config = load_t5(T5_SMALL)
t5 = t5.to(device).eval()
print(f"t5-small: {sum(p.numel() for p in t5.parameters()) / 1e6:.1f} million parameters")

t5-small: 60.5 million parameters

The check: the same input and the same decoder input through both models, teacher-forced, so every position is compared at once. The decoder’s first input is the padding token, T5’s start-of-output.

batch = tokenizer(prompts, return_tensors="pt", padding=True).to(device)
written = tokenizer(answers, return_tensors="pt", padding=True).input_ids.to(device)
# the padding token is also T5's start token
start = torch.zeros(len(prompts), 1, dtype=torch.long, device=device)
decoder_input = torch.cat([start, written[:, :-1]], dim=1)
with torch.no_grad():
    ours = t5.decode(decoder_input, t5.encode(batch.input_ids, batch.attention_mask),
                     batch.attention_mask)
    theirs = reference(**batch, decoder_input_ids=decoder_input).logits
print("largest difference in logits:", (ours - theirs).abs().max().item())

largest difference in logits: 3.814697265625e-05

Generation is lesson 4’s loop with an encoder in front: encode the input once, start the decoder from the padding token, and keep feeding back the most probable token until every row has written the end-of-text token </s>.

@torch.no_grad()
def generate(model, ids, mask, max_new_tokens=64, eos=1, pad=0, temperature=0.0,
             generator=None):
    memory = model.encode(ids, mask)
    caches = [{"self": {}, "cross": {}} for _ in model.decoder]
    token = torch.full((ids.shape[0], 1), pad, device=ids.device)
    written = []
    done = torch.zeros(ids.shape[0], dtype=torch.bool, device=ids.device)
    for step in range(max_new_tokens):
        logits = model.decode(token, memory, mask, caches, past=step)[:, -1].float()
        if temperature == 0:
            token = logits.argmax(dim=-1, keepdim=True)  # greedy
        else:
            probs = (logits / temperature).softmax(dim=-1)
            token = torch.multinomial(probs, 1, generator=generator)
        token = token.masked_fill(done[:, None], pad)  # finished rows write padding
        written.append(token)
        done = done | (token[:, 0] == eos)
        if done.all():
            break
    return torch.cat(written, dim=1)


ours = tokenizer.batch_decode(generate(t5, batch.input_ids, batch.attention_mask),
                              skip_special_tokens=True)
print("our generation matches Hugging Face's:", ours == answers)
for text in ours[:2]:
    print(" ", text)

our generation matches Hugging Face's: True
  Die Wachen kommen morgens um sieben Uhr im Museum an.
  Le train de nuit à Vienne quitte à 11 heures.

5. What the positions learned, and where the decoder looks

The 32 biases of each head are all T5 knows about word order, so they are worth a look.

Each row below is one of the eight heads of t5-small’s encoder; each column a distance from the query, from 20 tokens to the left to 20 to the right; the color is the bias the head adds to the score of a key at that distance.

def bias_by_distance(attention, distances):
    zero = torch.zeros(1, dtype=torch.long, device=device)
    with torch.no_grad():
        bias = attention.position_bias(zero, distances.to(device))
    return bias[0, :, 0].cpu()  # heads, distances


distances = torch.arange(-20, 21)
encoder_bias = bias_by_distance(t5.encoder[0].attention, distances)
limit = encoder_bias.abs().max().item()
# negative biases in teal, positive in rose, zero white
ends = [COLORS["enc"], COLORS["card"], COLORS["dec"]]
diverging = LinearSegmentedColormap.from_list("bias", ends)
fig, ax = plt.subplots(figsize=(6.4, 2.6))
image = ax.imshow(encoder_bias, cmap=diverging, vmin=-limit, vmax=limit, aspect="auto")
ax.set_xticks(range(0, 41, 5), [str(d) for d in range(-20, 21, 5)])
ax.set_yticks(range(8), [f"head {h}" for h in range(8)])
ax.set_xlabel("key position minus query position")
ax.grid(False)
fig.colorbar(image, ax=ax, label="bias added to the score")
plt.show()

Output of the code above

The heads have split the work. Some favour the left side and some the right; one rises sharply on the token just before the query, one on the token just after; one strongly avoids the query’s own position, so that it always reads a neighbour; and one barely cares about distance at all and attends by content.

A model that is told nothing about order beyond these 32 numbers per head still builds a sense of left, right and next door. Beyond about ten tokens the biases flatten out: the buckets there are wide, and far is simply far.

Cross-attention is where the encoder–decoder earns its name, and it can be watched. We translate the sentence of section 1 again with our model, keep the cross-attention weights of every decoder layer, and average the eight heads of one layer: each row is a German token being written, each column an English token of the input.

source = tokenizer([prompts[0]], return_tensors="pt").to(device)
german = generate(t5, source.input_ids, source.attention_mask)
start = torch.zeros(1, 1, dtype=torch.long, device=device)
decoder_input = torch.cat([start, german[:, :-1]], dim=1)
with torch.no_grad():
    memory = t5.encode(source.input_ids, source.attention_mask)
    _, cross = t5.decode(decoder_input, memory, source.attention_mask, keep_weights=True)
# SentencePiece marks a word's start with "▁"; we show it as a space
def readable(ids):
    return [t.replace("▁", " ") for t in tokenizer.convert_ids_to_tokens(ids)]


source_tokens, german_tokens = readable(source.input_ids[0]), readable(german[0])
LAYER = 4
weights = cross[LAYER][0].mean(dim=0).cpu()  # rows: German tokens, columns: English tokens
heatmap(weights, rows=german_tokens, cols=source_tokens, cell=0.3)
plt.show()

Output of the code above

Most German tokens find their English source: Wach looks at guard, kommen at arrive, sieben at seven, Museum at museum. German puts the particle of ankommen at the very end of the sentence, and the last word, an, looks back at arrive. Nobody gave the model an alignment; it learned one because translating needs it.

Much of the attention, though, goes to the task prefix, to and the colon, on nearly every row: a head that has nothing useful to read parks its weight on a token that carries no content, the same habit as the attention sinks of lesson 9.

6. Pre-training by repair: T5 and BART

An encoder–decoder needs a task with an input and an output, and raw text has neither. Both T5 and BART make one up the same way: damage the text, give the damaged version to the encoder, and train the decoder to undo the damage. They damage it differently.

Two ways to pre-train an encoder–decoder on raw text. T5's span corruption replaces spans of the sentence with sentinel tokens, extra_id_0 and extra_id_1, and the target lists only the dropped spans, each after its sentinel, so the target is short. BART's text infilling replaces each span with a single mask token, which can stand for any number of words, and the target is the whole original sentence Two ways to pre-train an encoder–decoder on raw text. T5's span corruption replaces spans of the sentence with sentinel tokens, extra_id_0 and extra_id_1, and the target lists only the dropped spans, each after its sentinel, so the target is short. BART's text infilling replaces each span with a single mask token, which can stand for any number of words, and the target is the whole original sentence

T5’s version is span corruption. It drops 15% of the tokens in spans of 3 on average, replaces each span with a sentinel token, <extra_id_0>, <extra_id_1> and so on, and asks the decoder for the dropped spans only, each after its sentinel, with one more sentinel at the end.

The targets are short, and that was the point: the choice was “made to reduce the computational cost of pre-training”. Across the many variants the paper tried, the details of the corruption mattered little. What mattered was that “denoising objectives outperformed language modeling”.

BART, which Facebook AI published the same month, is the same design seen from the other side: “a denoising autoencoder for pretraining sequence-to-sequence models”, whose bidirectional encoder and left-to-right decoder “can be seen as generalizing BERT […], GPT […], and other recent pretraining schemes” [5].

Its best corruption was text infilling: span lengths are drawn from a Poisson distribution with a mean of 3, and each span becomes a single mask token, even a span of length zero, so the model has to “predict how many tokens are missing from a span”.

The large model also shuffled the order of the sentences. BART’s decoder writes the whole document back, which made it a strong generator: on XSum, a summarisation dataset the BART paper calls highly abstractive, it beat the best earlier model by “roughly 3.5 points on all ROUGE metrics”.

Inside, BART is the Transformer of 2017 with GPT’s GELU in place of ReLU, which makes its layers closer to BERT’s than to T5’s: learned position embeddings for up to 1,024 positions, a LayerNorm with its mean and bias after each part of the layer, biases in every linear layer, and GPT-2’s byte-level vocabulary, in which < and every other character have a token.

The model we load below, bart-base, has 6 layers on each side and a width of 768.

T5’s recipe comes down to two numbers: 15% of the tokens are dropped, in spans of 3 tokens on average, at random places.

Google’s code first decides how many tokens to drop and how many spans to cut them into, then splits both the dropped and the kept tokens into that many pieces of random length, and interleaves them. Here it is on the example sentence from the T5 paper:

def random_split(total, parts, rng):
    # a random way to write `total` as a sum of `parts` positive whole numbers
    cuts = sorted(rng.sample(range(1, total), parts - 1))
    edges = [0] + cuts + [total]
    return [edges[i + 1] - edges[i] for i in range(parts)]


def span_corruption(tokens, rng, density=0.15, mean_span=3):
    n = len(tokens)
    dropped = min(max(round(n * density), 1), n - 1)
    spans = max(round(dropped / mean_span), 1)
    dropped_lengths = random_split(dropped, spans, rng)
    kept_lengths = random_split(n - dropped, spans, rng)
    source, target, i = [], [], 0
    for k in range(spans):
        source += tokens[i:i + kept_lengths[k]]
        i += kept_lengths[k]
        sentinel = f"<extra_id_{k}>"
        source.append(sentinel)
        target.append(sentinel)
        target += tokens[i:i + dropped_lengths[k]]
        i += dropped_lengths[k]
    target.append(f"<extra_id_{spans}>")  # the last sentinel says: nothing more
    return source, target


paper_sentence = tokenizer.tokenize("Thank you for inviting me to your party last week.")
paper_sentence = [t.replace("▁", " ") for t in paper_sentence]
source, target = span_corruption(paper_sentence, random.Random(3), density=0.3, mean_span=2)
token_chips([source, target], labels=["encoder reads", "decoder writes"],
            strong=lambda row, col, token: token.startswith("<extra"))
plt.show()

Output of the code above

The encoder sees the sentence with holes, each hole one sentinel however many tokens it hides. The decoder’s target lists what is missing, in order, and closes with one more sentinel.

A sentence of 11 tokens would lose only two at T5’s rate, so the cell doubles it; on a 512-token passage of C4 the real recipe cuts about two dozen holes, in every kind of place.

So does a pre-trained T5 repair text? Ask it to fill the spans the T5 paper drops. We use t5-base here, which we also need in section 8; the smallest model fills them less reliably.

T5_BASE = "google-t5/t5-base"
base_reference = T5ForConditionalGeneration.from_pretrained(T5_BASE).to(device).eval()
damaged = ["Thank you <extra_id_0> me to your party <extra_id_1> week.",
           "The guards <extra_id_0> at the museum at <extra_id_1> in the morning.",
           "The night train to <extra_id_0> leaves at eleven and arrives in <extra_id_1> "
           "the next morning."]
batch = tokenizer(damaged, return_tensors="pt", padding=True).to(device)
with torch.no_grad():
    repaired = base_reference.generate(**batch, max_new_tokens=16, num_beams=1,
                                       do_sample=False)
def fill_spans(damaged_text, written):
    # the output is "<extra_id_0> span <extra_id_1> span …"; each span goes back in place
    pieces = re.split(r"<extra_id_\d+>", written.replace("<pad>", "").replace("</s>", ""))
    filled = damaged_text
    for k in range(damaged_text.count("<extra_id_")):
        span = pieces[k + 1].strip() if k + 1 < len(pieces) else "?"
        filled = filled.replace(f"<extra_id_{k}>", f"[{span}]")
    return filled


rows = []
for text, ids in zip(damaged, repaired):
    written = tokenizer.decode(ids, skip_special_tokens=False)
    shown = written.replace("<pad>", "")[:48] + "…"
    rows.append({"damaged": text, "t5-base writes": shown,
                 "the spans put back": fill_spans(text, written)})
table(pd.DataFrame(rows))
damagedt5-base writesthe spans put back
Thank you <extra_id_0> me to your party <extra_id_1> week.<extra_id_0> for inviting<extra_id_1> this<extra…Thank you [for inviting] me to your party [this] week.
The guards <extra_id_0> at the museum at <extra_id_1> in the morning.<extra_id_0> arrive<extra_id_1> 8<extra_id_2> ar…The guards [arrive] at the museum at [8] in the morning.
The night train to <extra_id_0> leaves at eleven and arrives in <extra_id_1> the next morning.<extra_id_0> London<extra_id_1> London<extra_id_…The night train to [London] leaves at eleven and arrives in [London] the next morning.

t5-base fills the gaps sensibly: for inviting and this in the paper’s sentence, a verb and a time for the guards, a city for the train.

The middle column shows what it actually writes. It doesn’t stop after the last gap: the target of span corruption ends with one more sentinel, and this model keeps going past it, inventing more gaps and repeating itself, so we read only as many spans as there are holes.

Most likely the released t5-base, trained on long passages of C4 mixed with supervised tasks, never saw a single short sentence with two holes.

BART’s version is closer to BERT’s masked language model, and to what the decoder of lesson 4 does. Each span becomes one <mask>, whatever its length, and the decoder writes the whole text back, so it also has to decide how many words each mask hid:

BART = "facebook/bart-base"
bart_tokenizer = AutoTokenizer.from_pretrained(BART)
bart = BartForConditionalGeneration.from_pretrained(BART).to(device).eval()
masked = ["Thank you <mask> me to your party <mask> week.",
          "The guards <mask> at the museum at <mask> in the morning.",
          "The night train to <mask> leaves at eleven and arrives in <mask> "
          "the next morning."]
batch = bart_tokenizer(masked, return_tensors="pt", padding=True).to(device)
with torch.no_grad():
    repaired = bart.generate(**batch, max_new_tokens=32, num_beams=1, do_sample=False)
written = bart_tokenizer.batch_decode(repaired, skip_special_tokens=True)
table(pd.DataFrame({"masked": masked, "BART writes": written}))
maskedBART writes
Thank you <mask> me to your party <mask> week.Thank you for inviting me to your party this week.
The guards <mask> at the museum at <mask> in the morning.The guards were at the museum at about 2:30 in the morning.
The night train to <mask> leaves at eleven and arrives in <mask> the next morning.The night train to the airport leaves at eleven and arrives in New York the next morning.

BART writes each sentence back whole and decides the length of each gap itself: one mask becomes one word, another becomes two. For the night train it invents a whole route.

The repairs are fluent and plausible rather than right, and that is all a pre-training objective can ask for, since nothing in the damaged text says which words were removed. What the model learns from millions of such guesses is how text hangs together, which is what fine-tuning then puts to work.

7. Spider: questions about databases the model has never seen

Spider was built at Yale in 2018. Eleven students spent “a total of 1,000 man-hours” writing “10,181 questions and 5,693 unique complex SQL queries on 200 databases with multiple tables, covering 138 different domains” [6].

What made it hard was the split: the databases of the test set never appear in training, so a model has to read a schema it has never seen.

The best model in the paper got 12.4% of the test queries exactly right. In 2021 T5-3B with PICARD, a decoder that rejects any token that would make the query invalid, reached 75.1% execution accuracy on the test set.

PICARD’s authors noted that without it, “12% of the SQL queries generated by the T5-3B model on Spider’s development set result in an execution error” [7]. GPT-4 with careful prompting took the leaderboard to 85.3% and then 86.6% in 2023, and in February 2024 the authors closed it.

The Spider authors distribute the databases as SQLite files; the Hugging Face copy we use is a re-host of the official archive, pinned to one revision, under the same CC BY-SA 4.0 licence.

SPIDER = "HAL-9001/spider-databases"
SPIDER_REVISION = "4a01bbac6520cd35b216db9e1724e5e1ada60aa4"
archive = hf_hub_download(SPIDER, "spider_data.zip", repo_type="dataset",
                          revision=SPIDER_REVISION)
with zipfile.ZipFile(archive) as z:
    z.extractall(WORK_DIR)
SPIDER_DIR = WORK_DIR / "spider_data"

with open(SPIDER_DIR / "train_spider.json") as f:
    train_set = json.load(f)
with open(SPIDER_DIR / "dev.json") as f:
    dev_set = json.load(f)
with open(SPIDER_DIR / "tables.json") as f:
    schemas = {db["db_id"]: db for db in json.load(f)}
train_dbs = {e["db_id"] for e in train_set}
dev_dbs = {e["db_id"] for e in dev_set}
print(f"training: {len(train_set):,} questions on {len(train_dbs)} databases")
print(f"development: {len(dev_set):,} questions on {len(dev_dbs)} databases")
print("databases in both:", len(train_dbs & dev_dbs))

training: 7,000 questions on 140 databases
development: 1,034 questions on 20 databases
databases in both: 0

Not one database is shared between the two sets, as the last line shows. Here is the first development question, about concert_singer, and the table it asks about:

example = dev_set[0]
print("question:", example["question"])
print("query:   ", example["query"])
database = SPIDER_DIR / "database" / example["db_id"] / (example["db_id"] + ".sqlite")
with sqlite3.connect(database) as connection:
    singers = pd.read_sql_query("SELECT * FROM singer", connection)
table(singers)

question: How many singers do we have?
query:    SELECT count(*) FROM singer
Singer_IDNameCountrySong_NameSong_release_yearAgeIs_male
1Joe SharpNetherlandsYou199252F
2TimbalandUnited StatesDangerous200832T
3Justin BrownFranceHey Oh201329T
4Rose WhiteFranceSun200341F
5John NizinikFranceGentleman201443T
6Tribal KingFranceLove201625T

The model gets the question and the schema as one line of text: the database’s name, then every table with its columns, the names exactly as SQL has to write them.

Nothing else: no rows, no types, no keys. That keeps the input short, and it means the model has to guess, from names alone, that Singer_ID in singer_in_concert points to singer.

def schema_text(db_id, only=None):
    db = schemas[db_id]
    columns = {}
    for table_index, column in db["column_names_original"][1:]:  # the first is "*"
        columns.setdefault(table_index, []).append(column)
    parts = []
    for i, name in enumerate(db["table_names_original"]):
        if only is None or name.lower() in only:
            parts.append(name + " : " + " , ".join(columns.get(i, [])))
    return db_id + " | " + " | ".join(parts)


def model_input(example, only=None):
    return example["question"] + " | " + schema_text(example["db_id"], only)


def model_target(example):
    # Spider's queries have stray spaces and line breaks
    return " ".join(example["query"].split())


print(model_input(example))
print(model_target(example))
lengths = [len(tokenizer(model_input(e)).input_ids) for e in train_set]
print("input tokens: median", int(np.median(lengths)), "longest", max(lengths))
print("training inputs longer than 512 tokens:", sum(n > 512 for n in lengths))
dev_lengths = [len(tokenizer(model_input(e)).input_ids) for e in dev_set]
print("development inputs longer than 512 tokens:", sum(n > 512 for n in dev_lengths))

# a schema cut at 512 tokens would teach the model about tables it was never shown,
# so a long one keeps only the tables its query uses
shortened = 0
for e, n in zip(train_set, lengths):
    e["input"] = model_input(e)
    if n > 512:
        words = set(re.findall(r"[a-z_][a-z0-9_]*", e["query"].lower()))
        e["input"] = model_input(e, only=words)
        shortened += 1
for e in dev_set:
    e["input"] = model_input(e)
lengths = [len(tokenizer(e["input"]).input_ids) for e in train_set]
print("training schemas shortened to the tables of their query:", shortened)
print("training inputs still longer than 512 tokens:", sum(n > 512 for n in lengths))

How many singers do we have? | concert_singer | stadium : Stadium_ID , Location , Name , Capacity , Highest , Lowest , Average | singer : Singer_ID , Name , Country , Song_Name , Song_release_year , Age , Is_male | concert : concert_ID , concert_Name , Theme , Stadium_ID , Year | singer_in_concert : concert_ID , Singer_ID
SELECT count(*) FROM singer
input tokens: median 140 longest 1892
training inputs longer than 512 tokens: 323
development inputs longer than 512 tokens: 0
training schemas shortened to the tables of their query: 323
training inputs still longer than 512 tokens: 0

A few training databases have schemas too long for 512 tokens. Cutting them at 512 would teach the model to write about tables it was never shown, which is exactly the mistake we least want it to learn, and dropping their questions would lose some of the hardest joins in the set.

So for those questions alone the schema keeps only the tables that the reference query uses. Every development schema fits whole, so the evaluation never relies on this.

The keys are the costly omission. Spider’s tables.json lists which column points to which, and a join needs exactly that; but writing the keys in would push hundreds of training inputs and dozens of development inputs past 512 tokens, so we leave them out. Section 9 shows what that costs.

8. A T5 that writes SQL

Before training anything, one problem. Ask the tokenizer to read a query with a comparison:

query = "SELECT count(*) FROM singer WHERE Age < 30"
ids = tokenizer(query).input_ids
print(tokenizer.convert_ids_to_tokens(ids))
print("back to text:", tokenizer.decode(ids, skip_special_tokens=True))
with_less = sum("<" in model_target(e) for e in train_set)
print(f"training queries with '<': {with_less} of {len(train_set):,}")

['▁', 'SEL', 'ECT', '▁count', '(', '*', ')', '▁FROM', '▁singer', '▁W', 'HER', 'E', '▁Age', '▁', '<unk>', '▁30', '</s>']
back to text: SELECT count(*) FROM singer WHERE Age  30
training queries with '<': 320 of 7,000

T5 can’t write “less than”. Its vocabulary was learned for prose, and < isn’t in it, so the tokenizer turns it into <unk>, the unknown token, and decoding drops it.

The query loses its comparison without an error: Age < 30 becomes Age 30. A model fine-tuned on such targets would learn to write queries that can’t run, for every question that needs a “less than”.

The fix costs nothing. The embedding matrix has more rows than the tokenizer has tokens, because the vocabulary was rounded up to a multiple of 128, and the last rows have never been used. We add < to the tokenizer, it takes the first free id, and the model already has a row for it:

The rows of t5-small's embedding matrix, not to scale. Rows 0 to 31,999 are the SentencePiece vocabulary, from the padding, end and unknown tokens to the last piece. Rows 32,000 to 32,099 are the 100 sentinels, extra_id_99 down to extra_id_0. Rows 32,100 to 32,127 are 28 rows that no token uses, left over because the matrix was rounded up to a multiple of 128. The new token for less-than takes row 32,100 The rows of t5-small's embedding matrix, not to scale. Rows 0 to 31,999 are the SentencePiece vocabulary, from the padding, end and unknown tokens to the last piece. Rows 32,000 to 32,099 are the 100 sentinels, extra_id_99 down to extra_id_0. Rows 32,100 to 32,127 are 28 rows that no token uses, left over because the matrix was rounded up to a multiple of 128. The new token for less-than takes row 32,100

print("vocabulary rows in the model:", t5.embed.weight.shape[0])
print("tokens the tokenizer knows:  ", len(tokenizer))
sql_tokenizer = AutoTokenizer.from_pretrained(T5_SMALL)
sql_tokenizer.add_tokens(["<"])
print("id of the new token '<':     ", sql_tokenizer.convert_tokens_to_ids("<"))
ids = sql_tokenizer(query).input_ids
print(sql_tokenizer.convert_ids_to_tokens(ids))

vocabulary rows in the model: 32128
tokens the tokenizer knows:   32100
id of the new token '<':      32100
['▁', 'SEL', 'ECT', '▁count', '(', '*', ')', '▁FROM', '▁singer', '▁W', 'HER', 'E', '▁Age', '<', '▁30', '</s>']

The new token sits in a row the model never trained, so at first it is noise; fine-tuning will give it a meaning, because a few hundred training queries use it. The tokenizer is the same for every size of T5, so sql_tokenizer serves t5-base too.

Now the training. We fine-tune t5-base, the same architecture as t5-small with 12 layers on each side and a width of 768, loaded into our own T5 class; the transplant needs no change.

Each step takes a random batch of questions, runs the encoder on the inputs and the decoder on the queries shifted right by the start token, and minimises the cross-entropy of every query token, ignoring padding, exactly as in lesson 4.

del base_reference, bart  # the reference models have done their job
if device == "cuda":
    torch.cuda.empty_cache()
model, model_config = load_t5(T5_BASE)
model = model.to(device)
print(f"t5-base: {sum(p.numel() for p in model.parameters()) / 1e6:.0f} million parameters")


def make_batch(examples):
    inputs = sql_tokenizer([e["input"] for e in examples], max_length=512,
                           truncation=True, padding=True, return_tensors="pt")
    targets = sql_tokenizer([model_target(e) for e in examples], max_length=256,
                            truncation=True, padding=True, return_tensors="pt").input_ids
    start = torch.zeros(len(examples), 1, dtype=torch.long)  # the start token
    decoder_input = torch.cat([start, targets[:, :-1]], dim=1)
    labels = targets.masked_fill(targets == sql_tokenizer.pad_token_id, -100)
    return inputs.input_ids, inputs.attention_mask, decoder_input, labels


def loss_on(model, batch):
    ids, mask, decoder_input, labels = [t.to(device) for t in batch]
    logits = model.decode(decoder_input, model.encode(ids, mask), mask)
    return F.cross_entropy(logits.flatten(0, 1), labels.flatten(), ignore_index=-100)


held_out = random.Random(SEED).sample(dev_set, 256 if not SMOKE else 16)
with torch.no_grad():
    model.eval()
    before = loss_on(model, make_batch(held_out[:32])).item()
print(f"loss on development queries before training: {before:.2f}")

t5-base: 223 million parameters
loss on development queries before training: 4.94

The plan: 1,100 steps, about 20 minutes on a T4. Batches hold up to 16 questions of similar length and at most 4,096 input tokens, so that the longest schemas still fit in a T4’s memory.

The optimizer is AdamW at a constant learning rate after a short warm-up, with gradients clipped at 1. Every 200 steps we measure the loss on 256 held-out development questions.

We train in float32. T5 was trained on TPUs with bfloat16 activations, which have the range of float32, and in float16 its feed-forward outputs overflow: in 2021, “t5-small was the only T5 model that was working in fp16” [8]. A T4 has no fast bfloat16, so float32 it is.

STEPS = 3 if SMOKE else 1100  # about 20 minutes on a T4
BATCH = 16
LR = 3e-4
WARMUP = 100
optimizer = torch.optim.AdamW(model.parameters(), lr=LR, weight_decay=0.0)


TOKEN_BUDGET = 4096  # input tokens per batch: 16 questions of 256 tokens, or 8 of 512
input_tokens = {}
for e, n in zip(train_set, lengths):
    input_tokens[e["question"] + e["query"]] = min(n, 512)


def batches_by_length(examples, size, rng):
    shuffled = examples[:]
    rng.shuffle(shuffled)
    batches = []
    for i in range(0, len(shuffled), size * 50):
        pool = sorted(shuffled[i:i + size * 50],
                      key=lambda e: input_tokens[e["question"] + e["query"]])
        batch = []
        for e in pool:
            longest = input_tokens[e["question"] + e["query"]]  # the pool is sorted
            if len(batch) == size or (len(batch) + 1) * longest > TOKEN_BUDGET:
                batches.append(batch)
                batch = []
            batch.append(e)
        batches.append(batch)
    rng.shuffle(batches)  # similar lengths within a batch, random order of batches
    return batches


def dev_loss(model):
    model.eval()
    total = 0.0
    with torch.no_grad():
        for i in range(0, len(held_out), 32):
            chunk = held_out[i:i + 32]
            total += loss_on(model, make_batch(chunk)).item() * len(chunk)
    model.train()
    return total / len(held_out)


rng = random.Random(SEED)
history = []
model.train()
started, step = time.time(), 0
while step < STEPS:
    for examples in batches_by_length(train_set, BATCH, rng):
        if step == STEPS:
            break
        for group in optimizer.param_groups:
            group["lr"] = LR * min(1.0, (step + 1) / WARMUP)
        loss = loss_on(model, make_batch(examples))
        loss.backward()
        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
        optimizer.step()
        optimizer.zero_grad()
        if step % 200 == 0:
            history.append({"step": step, "minutes": (time.time() - started) / 60,
                            "train loss": loss.item(), "dev loss": dev_loss(model)})
        step += 1
history.append({"step": step, "minutes": (time.time() - started) / 60,
                "train loss": loss.item(), "dev loss": dev_loss(model)})
model.eval()
history = pd.DataFrame(history)
minutes = (time.time() - started) / 60
passes = step * BATCH / len(train_set)
print(f"{step:,} steps in {minutes:.0f} minutes, about {passes:.1f} passes over the data")

1,100 steps in 21 minutes, about 2.5 passes over the data
fig, ax = plt.subplots(figsize=(6.4, 2.8))
ax.plot(history["step"], history["train loss"], color=COLORS["enc"], label="training batch")
ax.plot(history["step"], history["dev loss"], color=COLORS["dec"],
        label="development questions")
ax.set_xlabel("step")
ax.set_ylabel("cross-entropy per token")
ax.set_ylim(0, None)
ax.legend(frameon=False)
plt.show()

Output of the code above

The loss falls fast: within the first 200 steps the development loss drops to a small fraction of where it started, and after that it barely moves while the training loss keeps going down. From there on the model gets better at the training questions than at SQL in general, the start of overfitting.

9. Checking by running

A loss says how probable the model finds the reference query, token by token. It doesn’t say whether the query the model writes is right: a query worded differently from the reference can be improbable and still return the right rows.

For SQL there is a better judge than comparing text: run it. Spider’s databases are real SQLite files, so we run the model’s query and the correct one on the same database and compare what comes back.

This is execution accuracy. It errs in both directions. Two correct queries can differ as strings, AGE > 34 and AGE >= 35, which is why exact matching underrates models.

And a wrong query can happen to return the right rows on one particular database, which is why Spider’s official metric since 2020 runs every query on a suite of databases built to tell such queries apart[9]. We use the one database Spider ships.

Two ways to grade a predicted query, on an illustrative table of four people aged 41, 29, 35 and 52. The reference is AGE greater than 34. A prediction of AGE at least 35 is a different string, so exact matching calls it wrong, but it returns the same three people and is right. A prediction of AGE greater than 30 also returns the same three people on this table, so execution calls it right although it is wrong; a second table with a person aged 32 tells the two apart, which is what a test suite does Two ways to grade a predicted query, on an illustrative table of four people aged 41, 29, 35 and 52. The reference is AGE greater than 34. A prediction of AGE at least 35 is a different string, so exact matching calls it wrong, but it returns the same three people and is right. A prediction of AGE greater than 30 also returns the same three people on this table, so execution calls it right although it is wrong; a second table with a person aged 32 tells the two apart, which is what a test suite does

Two results are the same if they hold the same rows. Order matters only if the correct query asks for one with ORDER BY; otherwise we compare the rows as sorted lists. A query that runs longer than five seconds is stopped and counts as wrong, and so does one that fails.

def run_query(db_id, query, seconds=5.0):
    path = SPIDER_DIR / "database" / db_id / (db_id + ".sqlite")
    connection = sqlite3.connect(path)
    # a few databases hold text that isn't valid UTF-8
    connection.text_factory = lambda raw: raw.decode(errors="replace")
    deadline = time.time() + seconds
    # SQLite calls the handler every 10,000 steps; returning True stops the query
    connection.set_progress_handler(lambda: time.time() > deadline, 10_000)
    try:
        return connection.execute(query).fetchall()
    finally:
        connection.close()


def same_result(db_id, predicted, gold):
    expected = run_query(db_id, gold)
    try:
        got = run_query(db_id, predicted)
    except Exception as error:
        return False, str(error)  # SQLite's message
    if "order by" in gold.lower():
        return got == expected, "ran"
    return sorted(map(repr, got)) == sorted(map(repr, expected)), "ran"


print(same_result("concert_singer", "SELECT count(*) FROM singer WHERE Age < 30",
                  "SELECT count(*) FROM singer WHERE age < 30"))
print(same_result("concert_singer", "SELECT count(*) FROM singer WHERE Age <= 30",
                  "SELECT count(*) FROM singer WHERE age < 30"))
print(same_result("concert_singer", "SELECT count(*) FROM singers",
                  "SELECT count(*) FROM singer WHERE age < 30"))

(True, 'ran')
(True, 'ran')
(False, 'no such table: singers')

Now every development question, 1,034 of them about 20 databases the model has never seen. Generation uses our own generate with the KV cache, in batches.

evaluated = dev_set if not SMOKE else dev_set[:24]
predictions = []
started = time.time()
for i in range(0, len(evaluated), 64):
    chunk = evaluated[i:i + 64]
    batch = sql_tokenizer([e["input"] for e in chunk], max_length=512, truncation=True,
                          padding=True, return_tensors="pt").to(device)
    written = generate(model, batch.input_ids, batch.attention_mask, max_new_tokens=192)
    predictions += sql_tokenizer.batch_decode(written, skip_special_tokens=True)
print(f"{len(predictions):,} queries written in {time.time() - started:.0f} s")

def normalized(query):
    query = " ".join(query.lower().replace('"', "'").split())
    return re.sub(r"\s*([,()<>=!])\s*", r"\1", query)  # spaces around , ( ) and < > = !


results = []
for e, predicted in zip(evaluated, predictions):
    correct, status = same_result(e["db_id"], predicted, model_target(e))
    results.append({"db_id": e["db_id"], "question": e["question"], "gold": model_target(e),
                    "predicted": predicted, "correct": correct, "status": status,
                    "as written": predicted == model_target(e),
                    "exact": normalized(predicted) == normalized(model_target(e)),
                    "no rows": len(run_query(e["db_id"], model_target(e))) == 0})
results = pd.DataFrame(results)
print(f"execution accuracy:                     {results['correct'].mean():.1%}")
print(f"the same string as the reference:       {results['as written'].mean():.1%}")
print(f"the same up to spaces, quotes and case: {results['exact'].mean():.1%}")
print("queries that failed to run:", (results["status"] != "ran").sum())
lucky = results["correct"] & results["no rows"] & ~results["exact"]
print("right only in that both return no rows:", lucky.sum())

1,034 queries written in 87 s
execution accuracy:                     42.6%
the same string as the reference:       11.2%
the same up to spaces, quotes and case: 32.0%
queries that failed to run: 406
right only in that both return no rows: 7

Just over two queries in five return the right rows, on twenty databases the model has never seen, after 20 minutes of fine-tuning a model of 220 million parameters.

The three ways of counting tell three different stories. Compared character for character, only about one query in nine matches the reference, and most of that gap is formatting: T5’s tokenizer drops the space before a comma, and Spider’s annotators mixed quotes and capital letters.

After normalising those, about a third match. Execution accepts roughly another tenth on top: queries written differently that mean the same thing.

A handful of those are right only by luck, because both queries return no rows at all on this database, the false positive of the figure above. And almost two queries in five don’t run.

Where does it go wrong? Spider’s own scripts sort queries into four levels of hardness by counting their parts; a rougher but transparent way is to look at what a correct query contains: a join, a nested SELECT, a GROUP BY, an ORDER BY, or none of these.

def query_parts(query):
    q = " " + query.lower() + " "
    return {"join": " join " in q, "nested select": q.count("select ") > 1,
            "group by": " group by " in q, "order by": " order by " in q}


rows = []
for name in ["join", "nested select", "group by", "order by", "none of these"]:
    if name == "none of these":
        chosen = [not any(query_parts(g).values()) for g in results["gold"]]
    else:
        chosen = [query_parts(g)[name] for g in results["gold"]]
    part = results[chosen]
    rows.append({"correct query has": name, "questions": len(part),
                 "execution accuracy": part["correct"].mean()})
parts = pd.DataFrame(rows)
table(parts)
correct query hasquestionsexecution accuracy
join4080.19
nested select1590.28
group by2770.39
order by2370.40
none of these3330.63
fig, ax = plt.subplots(figsize=(6.4, 2.6))
ax.barh(parts["correct query has"], parts["execution accuracy"], color=COLORS["enc"])
for k in range(len(parts)):
    accuracy, count = parts["execution accuracy"][k], parts["questions"][k]
    ax.text(accuracy + 0.01, k, f"{accuracy:.0%} of {count}", va="center",
            color=COLORS["muted"])
ax.invert_yaxis()
ax.set_xlim(0, 1)
ax.set_xlabel("execution accuracy")
plt.show()

Output of the code above

The pattern is the one Spider was built to expose. Queries with none of these parts, one table with a filter or a count, are right almost two times in three.

Every added clause costs, and joins cost the most: fewer than one in five is right. To join two tables, the model has to know which column of one points to which column of the other. Our input doesn’t say, so it guesses from the names.

Three kinds of answers are worth looking at one by one: queries that are right although they differ from the reference, queries that don’t run, and queries that run and return the wrong thing.

right_but_different = results[results["correct"] & ~results["exact"]]  # really different
failed = results[results["status"] != "ran"]
wrong = results[(results["status"] == "ran") & ~results["correct"]]
for title, frame in [("right, but not the reference string", right_but_different),
                     ("does not run", failed), ("runs, wrong result", wrong)]:
    print(f"== {title}: {len(frame)}")
    for row in frame.head(3).itertuples():
        print("  question: ", row.question)
        print("  reference:", row.gold)
        note = "" if row.status == "ran" else f"({row.status})"
        print("  model:    ", row.predicted, note)
        print()

== right, but not the reference string: 117
  question:  Show the stadium name and capacity with most number of concerts in year 2014 or after.
  reference: SELECT T2.name , T2.capacity FROM concert AS T1 JOIN stadium AS T2 ON T1.stadium_id = T2.stadium_id WHERE T1.year >= 2014 GROUP BY T2.stadium_id ORDER BY count(*) DESC LIMIT 1
  model:     SELECT T1.name, max(T1.capacity) FROM stadium AS T1 JOIN concert AS T2 ON T1.stadium_id = T2.stadium_id WHERE T2.year >= 2014 GROUP BY T2.stadium_id ORDER BY count(*) DESC LIMIT 1 

  question:  What is the name and capacity of the stadium with the most concerts after 2013 ?
  reference: select t2.name , t2.capacity from concert as t1 join stadium as t2 on t1.stadium_id = t2.stadium_id where t1.year > 2013 group by t2.stadium_id order by count(*) desc limit 1
  model:     SELECT T1.name, T1.capacity FROM stadium AS T1 JOIN concert AS T2 ON T1.stadium_id = T2.stadium_id WHERE T2.year > 2013 GROUP BY T1.name ORDER BY count(*) DESC LIMIT 1 

  question:  Show names for all stadiums except for stadiums having a concert in year 2014.
  reference: SELECT name FROM stadium EXCEPT SELECT T2.name FROM concert AS T1 JOIN stadium AS T2 ON T1.stadium_id = T2.stadium_id WHERE T1.year = 2014
  model:     SELECT name FROM stadium EXCEPT SELECT T1.name FROM stadium AS T1 JOIN concert AS T2 ON T1.stadium_id = T2.stadium_id WHERE T2.year = 2014 

== does not run: 406
  question:  Show the name and the release year of the song by the youngest singer.
  reference: SELECT song_name , song_release_year FROM singer ORDER BY age LIMIT 1
  model:     SELECT T1.name, T1.song_release_year FROM singer AS T1 JOIN concert AS T2 ON T1.singer_id = T2.singer_id ORDER BY T2.age LIMIT 1 (no such column: T2.singer_id)

  question:  What are the names and release years for all the songs of the youngest singer?
  reference: SELECT song_name , song_release_year FROM singer ORDER BY age LIMIT 1
  model:     SELECT T1.name, T1.song_release_year FROM singer AS T1 JOIN concert AS T2 ON T1.singer_id = T2.singer_id ORDER BY T2.age LIMIT 1 (no such column: T2.singer_id)

  question:  Show location and name for all stadiums with a capacity between 5000 and 10000.
  reference: SELECT LOCATION , name FROM stadium WHERE capacity BETWEEN 5000 AND 10000
  model:     SELECT T1.location, T1.name FROM stadium AS T1 JOIN concert AS T2 ON T1.stadium_id = T2.stadium_id WHERE T2.capacity BETWEEN 5000 AND 10000 (no such column: T2.capacity)

== runs, wrong result: 187
  question:  Show name, country, age for all singers ordered by age from the oldest to the youngest.
  reference: SELECT name , country , age FROM singer ORDER BY age DESC
  model:     SELECT name, country FROM singer ORDER BY age 

  question:  How many singers are from each country?
  reference: SELECT country , count(*) FROM singer GROUP BY country
  model:     SELECT count(*), Country FROM singer GROUP BY Country 

  question:  What is the maximum capacity and the average of all stadiums ?
  reference: select max(capacity), average from stadium
  model:     SELECT max(capacity), avg(capacity) FROM stadium

The first group is why execution accuracy exists: the model groups by the name instead of the id, joins the tables from the other side, or writes an aggregate where the reference doesn’t need one, and the rows come out the same. Some of these are right on this database only; a suite of test databases would catch a few.

The second group is the model’s main failure, and it isn’t about SQL. The queries are well formed, but they reach for a column or a table that doesn’t exist, often through a join the question didn’t need.

The third group is where our own check shows its limits. “How many singers are from each country?” answered with SELECT count(*), Country is right: it holds the same numbers in the other column order, and our comparison, which is strict about the order of columns, counts it as wrong.

The last example is a matter of annotation: the reference reads “the average” as a column called Average, the model as avg(capacity), and a person could read the question either way. No metric of text-to-SQL is free of such cases; what matters is knowing which way each one errs.

The queries that don’t run, sorted by what SQLite said about them:

def error_kind(message):
    for kind in ["no such column", "syntax error", "no such table", "ambiguous column",
                 "interrupted"]:
        if kind in message:
            return kind
    return "other"


kinds = Counter(error_kind(message) for message in failed["status"])
table(pd.DataFrame(kinds.most_common(), columns=["SQLite says", "queries"]))
SQLite saysqueries
no such column343
no such table32
syntax error19
ambiguous column7
other5

Five failures in six are of one kind: the query names a column that the table doesn’t have. Syntax errors are rare. After 20 minutes the model writes grammatical SQL; what it can’t do reliably is read the schema.

The text-to-SQL literature calls this schema linking, and it is why the best T5 systems on Spider gave the model some of each database’s contents as well as its names. Here is one failure drawn against the real schema:

invented = failed[failed["status"].str.startswith("no such column")]
if len(invented) > 0:
    case = invented.iloc[0]
    db = schemas[case["db_id"]]
    print("question:", case["question"])
    print("model:   ", case["predicted"])
    print("SQLite:  ", case["status"])
    rows, labels = [], []
    for i, name in enumerate(db["table_names_original"]):
        rows.append([c for t, c in db["column_names_original"][1:] if t == i])
        labels.append(name)
    rows.append([case["status"].split(": ", 1)[1]])
    labels.append("the model wrote")
    token_chips(rows, labels=labels, mark_spaces=False,
                strong=lambda row, col, token: row == len(rows) - 1)
    plt.show()

question: Show the name and the release year of the song by the youngest singer.
model:    SELECT T1.name, T1.song_release_year FROM singer AS T1 JOIN concert AS T2 ON T1.singer_id = T2.singer_id ORDER BY T2.age LIMIT 1
SQLite:   no such column: T2.singer_id

Output of the code above

The model knows that a singer has an id and that concerts involve singers, so it joins singer to concert on singer_id. But concert has no such column: singers and concerts meet in a third table, singer_in_concert, and nothing in our input says so.

The foreign keys we left out would have. The question didn’t even need a join: the youngest singer’s song is in singer itself.

10. Letting SQLite choose

A query that doesn’t run is the cheapest mistake to catch: SQLite says so in a millisecond, without knowing the right answer. PICARD used this inside the decoder, rejecting any token that would make the query unparseable.

Earlier, Wang et al. ran partly written queries during beam search and dropped the beams whose execution failed, which they called execution-guided decoding[10].

We do the simplest version of the idea: sample four more queries per question at a temperature of 0.7, put the greedy query first, and keep the first one that SQLite runs without an error. The choice looks only at the database, never at the reference answer, so it is a fair test.

The last line, at least one of five is right, is not a method: it peeks at the reference. It shows how much a perfect judge could get out of the same five candidates.

SAMPLES = 4
generator = torch.Generator(device=device).manual_seed(SEED)
candidates = [[query] for query in predictions]  # the greedy query comes first
started = time.time()
for i in range(0, len(evaluated), 64):
    chunk = evaluated[i:i + 64]
    batch = sql_tokenizer([e["input"] for e in chunk], max_length=512, truncation=True,
                          padding=True, return_tensors="pt").to(device)
    for _ in range(SAMPLES):
        written = generate(model, batch.input_ids, batch.attention_mask, max_new_tokens=192,
                           temperature=0.7, generator=generator)
        queries = sql_tokenizer.batch_decode(written, skip_special_tokens=True)
        for j, query in enumerate(queries):
            candidates[i + j].append(query)
print(f"{SAMPLES} more queries per question in {time.time() - started:.0f} s")


def first_that_runs(db_id, queries):
    for query in queries:
        try:
            run_query(db_id, query)
            return query
        except Exception:
            continue
    return queries[0]  # none runs: keep the greedy one


chosen, any_right = [], []
for e, queries in zip(evaluated, candidates):
    chosen.append(first_that_runs(e["db_id"], queries))
    any_right.append(any(same_result(e["db_id"], q, model_target(e))[0] for q in queries))
results["chosen"] = chosen
results["chosen correct"] = [same_result(e["db_id"], q, model_target(e))[0]
                             for e, q in zip(evaluated, chosen)]
results["any of 5 correct"] = any_right
still_failing = 0
for e, q in zip(evaluated, chosen):
    try:
        run_query(e["db_id"], q)
    except Exception:
        still_failing += 1
print(f"greedy:                      {results['correct'].mean():.1%}")
print(f"first of 5 that runs:        {results['chosen correct'].mean():.1%}")
print(f"at least one of 5 is right:  {results['any of 5 correct'].mean():.1%}")
print("queries that still fail to run:", still_failing)

4 more queries per question in 348 s
greedy:                      42.6%
first of 5 that runs:        48.8%
at least one of 5 is right:  52.8%
queries that still fail to run: 260
labels = ["greedy", "first of 5\nthat runs", "any of 5\nis right"]
values = [results["correct"].mean(), results["chosen correct"].mean(),
          results["any of 5 correct"].mean()]
fig, ax = plt.subplots(figsize=(6.4, 2.4))
ax.barh(labels, values, color=[COLORS["enc"], COLORS["enc"], COLORS["line"]])
for k, value in enumerate(values):
    ax.text(value + 0.01, k, f"{value:.1%}", va="center", color=COLORS["muted"])
ax.invert_yaxis()
ax.set_xlim(0, 1)
ax.set_xlabel("execution accuracy on 1,034 development questions")
plt.show()

Output of the code above

Letting SQLite reject the queries that fail adds about six points of execution accuracy, with no new training and nothing but the database’s word that a query runs.

More than a third of the failing questions get a working replacement from the four samples. The last bar shows what is left for a better judge: a perfect choice among the same five candidates would add only about four points more.

For nearly half of the questions, none of the five candidates is right. Choosing better won’t fix those; a bigger model, longer training, or the schema’s keys and some of its values in the input would.

11. Why decoder-only models won anyway

T5’s own comparison put the encoder–decoder first: “For all tasks, the encoder-decoder architecture with the denoising objective performed best.” Three years later the largest models were all decoders. The reasons are worth knowing, because none of them is that the encoder–decoder is worse.

The first is what got measured. T5 measured transfer learning: pre-train, then fine-tune for each task. The models that changed the field were used another way, prompted with no fine-tuning at all.

BigScience tested exactly this before choosing the architecture of BLOOM, training models of 5 billion parameters on 170 billion tokens, and found that “the best objective and architecture is the opposite in these two settings”: a causal decoder trained on plain language modelling did best right after pre-training, an encoder–decoder trained on masked language modelling did best once multi-task fine-tuning was added[11].

BLOOM’s report [12] stated the market’s verdict: “all state-of-the-art language models over 100 billion parameters are causal decoder-only models”.

The second is memory. Each token passes through only one stack, so an encoder–decoder costs about as much compute as a decoder-only model with half its parameters. But it has “twice the parameters of the decoder-only model, and hence twice the memory footprint”, and when the limit is what fits on the GPUs, that counts.

The third is how much the objective teaches per token: denoising learns from the 15% of tokens it hides, language modelling from nearly all of them, which Yi Tay, the first author of UL2, calls less “loss exposure” [13].

What an encoder–decoder costs against a decoder-only model. The decoder-only model is one stack with P parameters, and every token, of the input and of the output, passes through it. The encoder–decoder has two stacks of P parameters each: input tokens pass through the encoder, output tokens through the decoder. Each token meets P parameters in both cases, so the compute is about the same, but the encoder–decoder has to keep twice the parameters in memory What an encoder–decoder costs against a decoder-only model. The decoder-only model is one stack with P parameters, and every token, of the input and of the output, passes through it. The encoder–decoder has two stacks of P parameters each: input tokens pass through the encoder, output tokens through the decoder. Each token meets P parameters in both cases, so the compute is about the same, but the encoder–decoder has to keep twice the parameters in memory

There is also an engineering argument, which is ours rather than a paper’s: one stack is simpler. A chat is one growing sequence in which every answer becomes part of the next input.

A decoder reads it with one KV cache; an encoder–decoder has to decide what counts as input and re-encode it as the conversation grows. Kernels, serving systems and fine-tuning tools were then built for the one shape that won, which made it cheaper still.

The verdict is not final. In 2025 Google turned pre-trained Gemma 2 decoders into encoder–decoders and reported that “under similar inference budget, encoder-decoder LLMs achieve comparable (often better) pretraining performance but substantially better finetuning performance than their decoder-only counterpart” [14].

Google’s announcement adds that a model with a 9B encoder and a 2B decoder runs at nearly the latency of the 2B model [15]. T5Gemma 2 followed in December 2025, multimodal, in sizes from 270M-270M to 4B-4B.

12. Where the encoder–decoder is the right tool today

Wherever the input is long or of a different kind from the output, and the output is short. Speech recognition: Whisper is an encoder–decoder, chosen because “this architecture has been well validated to scale reliably” [16].

Translation: Meta’s NLLB-200 and the Helsinki-NLP MarianMT models. Image generation: the frozen encoder half of a T5-XXL turns the prompt into vectors in Imagen, Stable Diffusion 3 and FLUX.1.

Imagen’s authors found “scaling the text encoder size to be significantly more impactful than the U-Net size”[17], and Stable Diffusion 3’s that dropping T5 costs little, except when the prompt asks for written text in the picture [18].

And structured jobs like ours: a question to SQL, a document to fields, a ticket to a category with an explanation. For such a job, a small encoder–decoder fine-tuned on a few thousand examples is cheap to run and, in our experience, hard to beat at its size.

What it is not for is open conversation and learning from examples in the prompt. The models people chat with are decoders, and lesson 6 is about what changed inside them since GPT-2.

T5’s passport, and the family tree with this lesson’s branch:

T5's passport. Published October 2019 by Google. Type: encoder–decoder. Pre-training: span corruption on C4, and in version 1.0 a mixture of supervised tasks. Attention: bidirectional in the encoder, causal self-attention and cross-attention in the decoder. Positions: no position embeddings, a learned bias per head over 32 log-spaced buckets of distance. Normalisation: pre-norm with scale only, no mean and no bias. Activation: ReLU in version 1.0, gated GELU from version 1.1. Vocabulary: SentencePiece, 32,128 rows, without the characters less-than and curly braces. Sizes: 60 million, 220 million, 770 million, 3 and 11 billion parameters. Good at: input-to-output tasks such as translation, extraction, SQL, and as a text encoder. Not for: open-ended chat, long free generation, in-context learning T5's passport. Published October 2019 by Google. Type: encoder–decoder. Pre-training: span corruption on C4, and in version 1.0 a mixture of supervised tasks. Attention: bidirectional in the encoder, causal self-attention and cross-attention in the decoder. Positions: no position embeddings, a learned bias per head over 32 log-spaced buckets of distance. Normalisation: pre-norm with scale only, no mean and no bias. Activation: ReLU in version 1.0, gated GELU from version 1.1. Vocabulary: SentencePiece, 32,128 rows, without the characters less-than and curly braces. Sizes: 60 million, 220 million, 770 million, 3 and 11 billion parameters. Good at: input-to-output tasks such as translation, extraction, SQL, and as a text encoder. Not for: open-ended chat, long free generation, in-context learning

The family of the Transformer as this course covers it, with lesson 5, T5 and BART, highlighted and lessons 1 to 4 done The family of the Transformer as this course covers it, with lesson 5, T5 and BART, highlighted and lessons 1 to 4 done

References

  1. Green, B. F., Wolf, A. K., Chomsky, C., & Laughery, K. (1961). Baseball: an automatic question-answerer. Western Joint IRE-AIEE-ACM Computer Conference, 219–224. ↩
  2. Lei, F., Chen, J., Ye, Y., et al. (2025). Spider 2.0: evaluating language models on real-world enterprise text-to-SQL workflows. ICLR 2025. ↩
  3. Raffel, C., Shazeer, N., Roberts, A., Lee, K., Narang, S., Matena, M., Zhou, Y., Li, W., & Liu, P. J. (2020). Exploring the limits of transfer learning with a unified text-to-text transformer. Journal of Machine Learning Research, 21(140), 1–67. ↩
  4. Zhang, B., & Sennrich, R. (2019). Root mean square layer normalization. NeurIPS 2019. ↩
  5. Lewis, M., Liu, Y., Goyal, N., Ghazvininejad, M., Mohamed, A., Levy, O., Stoyanov, V., & Zettlemoyer, L. (2020). BART: denoising sequence-to-sequence pre-training for natural language generation, translation, and comprehension. ACL 2020. ↩
  6. Yu, T., Zhang, R., Yang, K., et al. (2018). Spider: a large-scale human-labeled dataset for complex and cross-domain semantic parsing and text-to-SQL task. EMNLP 2018. Data: yale-lily.github.io/spider, CC BY-SA 4.0. ↩
  7. Scholak, T., Schucher, N., & Bahdanau, D. (2021). PICARD: parsing incrementally for constrained auto-regressive decoding from language models. EMNLP 2021. ↩
  8. Hugging Face forum (2021). T5 fp16 issue is fixed. ↩
  9. Zhong, R., Yu, T., & Klein, D. (2020). Semantic evaluation for text-to-SQL with distilled test suites. EMNLP 2020. ↩
  10. Wang, C., Tatwawadi, K., Brockschmidt, M., Huang, P.-S., Mao, Y., Polozov, O., & Singh, R. (2018). Robust text-to-SQL generation with execution-guided decoding. arXiv:1807.03100. ↩
  11. Wang, T., Roberts, A., Hesslow, D., Le Scao, T., et al. (2022). What language model architecture and pretraining objective work best for zero-shot generalization? ICML 2022. ↩
  12. BigScience Workshop (2022). BLOOM: a 176B-parameter open-access multilingual language model. arXiv:2211.05100. ↩
  13. Tay, Y. (2024). What happened to BERT & T5? On Transformer encoders, PrefixLM and denoising objectives. ↩
  14. Zhang, B., Moiseev, F., Ainslie, J., et al. (2025). Encoder-decoder Gemma: improving the quality-efficiency trade-off via adaptation. arXiv:2504.06225. ↩
  15. Google (2025). T5Gemma: a new collection of encoder-decoder Gemma models. Google Developers Blog. ↩
  16. Radford, A., Kim, J. W., Xu, T., Brockman, G., McLeavey, C., & Sutskever, I. (2022). Robust speech recognition via large-scale weak supervision. arXiv:2212.04356. ↩
  17. Saharia, C., Chan, W., Saxena, S., et al. (2022). Photorealistic text-to-image diffusion models with deep language understanding. NeurIPS 2022. ↩
  18. Esser, P., Kulal, S., Blattmann, A., et al. (2024). Scaling rectified flow transformers for high-resolution image synthesis. ICML 2024. ↩
  19. Tay, Y., Dehghani, M., Tran, V. Q., et al. (2023). UL2: unifying language learning paradigms. ICLR 2023.
  20. Zhang, B., Suganthan, P., Liu, G., et al. (2025). T5Gemma 2: seeing, reading, and understanding longer. arXiv:2512.14856.
  21. NLLB Team (2022). No language left behind: scaling human-centered machine translation. arXiv:2207.04672.
]]> Eva Koroleva GPT from scratch: GPT-2, the KV cache and a model that learns chess 2026-04-28T13:00:00+00:00 2026-11-03T00:00:00+00:00 https://blog.korolevalab.com/blog/2026/transformer-architectures-04-gpt Transformer architectures: from BERT to diffusion · lesson 4 of 15 · all lessons · run on Kaggle · open in Colab · code on GitHub

In September 2023 OpenAI released gpt-3.5-turbo-instruct, an unremarkable model for completing text, and within days people noticed that it played chess.

Nicholas Carlini gave it the moves of a game so far, in PGN, the notation of chess books and databases, and asked it for the next one: “GPT-4, for example, can’t even play a full game against me without making a few illegal moves. GPT-3.5-turbo-instruct can beat me” [1].

Against the bots of Lichess it later rated 1755, the level of a good amateur[2]. A year on, Dynomight tested a dozen other language models, open ones and OpenAI’s newer chat models among them, and rated every one of them “Terrible”[3].

Part of the explanation sits in an appendix of an OpenAI paper: GPT-4’s pre-training data “included chess games in the format of move sequence known as Portable Game Notation (PGN)”, and only games of players rated 1800 or higher [4].

Nobody taught these models the rules. They read game records, and a game record is a strange text: it lists moves and never shows a position.

To write a legal move you have to know where the pieces stand, and nothing in the text says. So does a model that only ever predicted the next character of game records keep a board somewhere inside it? That question fits on a free GPU.

A chess GPT reads a game record one character at a time, a semicolon, 1, a dot, e, 4, a space, e, 5 and so on, and predicts the next character. Nothing in the input says where the pieces stand. A linear probe, one matrix, reads the residual stream after a middle layer and recovers what stands on each of the 64 squares A chess GPT reads a game record one character at a time, a semicolon, 1, a dot, e, 4, a space, e, 5 and so on, and predicts the next character. Nothing in the input says where the pieces stand. A linear probe, one matrix, reads the residual stream after a middle layer and recovers what stands on each of the 64 squares

All of these models are GPTs, descendants of a 2018 OpenAI model trained to do one thing: predict the next token. In this lesson you will write GPT-2 yourself, load OpenAI’s weights into your code, and watch it generate and change its mind layer by layer.

Then you will train a small GPT of your own on Lichess games for 25 minutes, see how well it plays, and look inside it for the board.

What you’ll do

  1. See what a decoder-only model is trained to do, and why one rule turned out to be enough.
  2. Read the blueprints of GPT-1 and the four GPT-2s from their configs, and count where the parameters live.
  3. Write GPT-2 in PyTorch, load OpenAI’s weights into it and match Hugging Face’s outputs.
  4. Generate text with a KV cache and measure what the cache saves.
  5. Watch GPT-2 change its mind layer by layer with the logit lens.
  6. Ask GPT-2 to play chess, then train a GPT of your own on Lichess games and see how often its moves are legal.
  7. Open the chess model up and find the board inside it with a linear probe.

What you need: a GPU. On Kaggle, Settings → Accelerator → GPU; on Colab, Runtime → Change runtime type → T4 GPU. The notebook downloads GPT-2 (550 MB) and a gigabyte of Lichess games, and trains a chess model for 25 minutes; on a T4 it runs in about 40 minutes.

Setup

One small helper package, tfs, draws the course’s plots and token chips and knows where the notebook runs (Kaggle, Colab or your own machine). Everything the lesson teaches is written out in the notebook itself.

import importlib.util
import subprocess
import sys

TFS = "git+https://github.com/xmarva/transformers-from-scratch@main"
if importlib.util.find_spec("tfs") is None:
    subprocess.check_call([sys.executable, "-m", "pip", "install", "-q", TFS])

from tfs import env
from tfs.viz import table, token_chips, use_course_style

env.quiet()  # no progress bars or environment warnings in the outputs
SEED = 42
env.seed_everything(SEED)
COLORS = use_course_style()
WORK_DIR = env.work_dir()

import importlib.util
import json
import math
import random
import re
import subprocess
import sys
import time
from collections import Counter

# python-chess knows the rules of chess; our models will never be told them
if importlib.util.find_spec("chess") is None:
    subprocess.check_call([sys.executable, "-m", "pip", "install", "-q", "chess==1.11.2"])
import chess
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import torch
import torch.nn as nn
import torch.nn.functional as F
from huggingface_hub import get_safetensors_metadata, hf_hub_download
from IPython.display import display
from safetensors.torch import load_file
from tfs.viz import heatmap
from transformers import AutoTokenizer
# the reference implementation, for the check
from transformers import GPT2LMHeadModel

SMOKE = env.smoke()  # CI runs the lesson on a small slice; you run it in full
device = "cuda" if torch.cuda.is_available() else "cpu"
assert device == "cuda" or SMOKE, "This lesson trains a model for 25 minutes: use a GPU."
if device == "cuda":
    print("GPU:", torch.cuda.get_device_name(0))

GPU: Tesla T4

1. One rule: predict the next token

BERT, in the first lesson, read a sentence with holes and filled them, looking both ways.

A GPT reads a text from left to right and, at every position, guesses the next token. That is the whole training task. The probability of a text \(x_1, \dots, x_T\) is written as a chain of such guesses,

\[p(x_1, \dots, x_T) = \prod_{t=1}^{T} p(x_t \mid x_1, \dots, x_{t-1}),\]

which is not an approximation but the chain rule of probability: any distribution over sequences can be written this way. The model’s job is to supply each factor. Training minimises the average negative log of the probability it gave to the token that actually came next:

\[\mathcal{L} = -\frac{1}{T} \sum_{t=1}^{T} \log p_\theta(x_t \mid x_{<t}).\]

Here \(\theta\) are the weights, and \(x_{<t}\) is everything before position \(t\). A loss of \(\mathcal{L}\) means the model gave the right token a probability of \(e^{-\mathcal{L}}\), on average in the geometric sense.

Two things make this rule cheap to train. The targets are the text itself shifted by one position, so any text is training data and nobody has to label anything.

And one pass over a sequence of \(T\) tokens trains all \(T\) predictions at once: the causal mask lets position \(t\) see positions \(1\) to \(t\) and nothing after, so the model can’t peek at the answer it is asked for.

Next-token prediction. The input is the text, the targets are the same text shifted one position left. The causal mask is a lower-triangular grid: each position attends to itself and every earlier position, never to later ones, so one forward pass makes a prediction at every position, each from the left context only Next-token prediction. The input is the text, the targets are the same text shifted one position left. The causal mask is a lower-triangular grid: each position attends to itself and every earlier position, never to later ones, so one forward pass makes a prediction at every position, each from the left context only

The rule has a second life after training. To generate, the model predicts a distribution over the next token, one token is picked from it and appended, and the longer text goes back in. Everything a chatbot writes is made this way, one token at a time.

2. Three GPTs

OpenAI’s first GPT, in June 2018, was a 12-layer Transformer decoder trained on BooksCorpus, “over 7,000 unique unpublished books”, chosen because “it contains long stretches of contiguous text” [5].

Pre-training took “1 month on 8 GPUs”; then a copy of the model was fine-tuned for each task, and it improved the state of the art “in 9 out of the 12 tasks studied”.

Four months later BERT beat it on GLUE by reading in both directions, and for two years the encoder looked like the winner.

GPT-2, in February 2019, was “a direct scale-up of GPT, with more than 10X the parameters and trained on more than 10X the amount of data”[6].

The data, WebText, was the text of pages linked from Reddit posts with at least 3 karma, which the paper calls “a heuristic indicator for whether other users found the link interesting, educational, or just funny”: 8 million documents, 40 GB.

The largest model, 1.5 billion parameters, was not fine-tuned for anything, and still set the state of the art on 7 of 8 language-modelling benchmarks; it even summarised articles when the prompt ended with “TL;DR:” [7].

OpenAI held it back: “Due to our concerns about malicious applications of the technology, we are not releasing the trained model.”

Newspapers turned that into a model “too dangerous to release”, a phrase OpenAI never used, and the weights came out in four stages, from the smallest in February to the full model in November 2019.

GPT-3, in May 2020, kept GPT-2’s design, with attention layers that alternate between full and locally banded patterns, and grew it to 175 billion parameters in 96 layers, trained on 300 billion tokens [8].

Its paper named in-context learning: show the model a few examples of a task in the prompt, with no change to its weights, and it does the task.

People asked to tell its news articles from real ones were right “∼ 52%” of the time, barely above chance. Its weights were never published. The chat models that followed, from ChatGPT to Llama, Qwen and DeepSeek, are all of this line: a stack of decoder layers predicting the next token.

The configs of GPT-1 and of the four GPT-2s are on the Hugging Face Hub, and the safetensors files say how many numbers they hold without downloading them:

GPT_FAMILY = ["openai-gpt", "gpt2", "gpt2-medium", "gpt2-large", "gpt2-xl"]
rows = []
for name in GPT_FAMILY:
    repo = "openai-community/" + name
    with open(hf_hub_download(repo, "config.json")) as f:
        c = json.load(f)
    stored = sum(get_safetensors_metadata(repo).parameter_count.values())
    rows.append({"model": name, "layers": c["n_layer"], "width": c["n_embd"],
                 "heads": c["n_head"], "context": c["n_positions"],
                 "vocabulary": c["vocab_size"], "numbers stored": stored})
family = pd.DataFrame(rows)
table(family)
modellayerswidthheadscontextvocabularynumbers stored
openai-gpt127681251240,478119,680,512
gpt212768121,02450,257137,022,720
gpt2-medium241,024161,02450,257379,988,992
gpt2-large361,280201,02450,257811,778,816
gpt2-xl481,600251,02450,2571,607,942,848

The width \(d\), the number of layers \(L\) and the vocabulary \(V\) fix the size of the model. Each layer holds four matrices of attention (\(Q\), \(K\), \(V\) and the output projection, \(4d^2\) weights) and two of the feed-forward network, \(d \times 4d\) and \(4d \times d\) (\(8d^2\)). With biases and two LayerNorms that is \(12d^2 + 13d\) per layer. The embeddings add \(Vd\) for tokens and \(Cd\) for positions, where \(C\) is the context, and the output layer adds nothing, for a reason we’ll see in a moment. Let’s check the formula against the files:

def gpt_parameters(layers, d, vocab, context):
    per_layer = 12 * d * d + 13 * d
    embeddings = vocab * d + context * d
    final_norm = 2 * d  # GPT-2 has one more LayerNorm at the end
    return layers * per_layer + embeddings + final_norm


family["by the formula"] = [gpt_parameters(r.layers, r.width, r.vocabulary, r.context)
                            for r in family.itertuples()]
family["gap"] = family["numbers stored"] - family["by the formula"]
family["one mask per layer"] = family["layers"] * family["context"] ** 2
table(family[["model", "numbers stored", "by the formula", "gap", "one mask per layer"]])
modelnumbers storedby the formulagapone mask per layer
openai-gpt119,680,512116,536,3203,144,1923,145,728
gpt2137,022,720124,439,80812,582,91212,582,912
gpt2-medium379,988,992354,823,16825,165,82425,165,824
gpt2-large811,778,816774,030,08037,748,73637,748,736
gpt2-xl1,607,942,8481,557,611,20050,331,64850,331,648

The formula is right, up to a gap that has a funny explanation. OpenAI’s checkpoints store the causal mask, a \(C \times C\) triangle of ones and zeros, as a tensor in every layer, so the files hold a million numbers per layer that aren’t weights. GPT-1’s gap is \(2d\) short of its masks: it has no final LayerNorm, which the formula counted.

The smallest GPT-2 has 124 million parameters. Its paper says 117M, and the larger ones 345M and 762M; OpenAI later admitted in its repository that “our original parameter counts were wrong” (openai/gpt-2). 117M is the size of GPT-1, which the smallest GPT-2 was meant to match: the extra seven million are the bigger vocabulary and the longer context.

Where those parameters live depends on the size. In the smallest GPT-2 the token embeddings are almost a third of the model; in GPT-2 XL they are a few percent, and the layers are everything:

fig, ax = plt.subplots(figsize=(6.4, 2.6))
names = family["model"][1:]
for k, r in enumerate(family[1:].itertuples()):
    embed = (r.vocabulary + r.context) * r.width / 1e6
    layers_part = r.layers * (12 * r.width ** 2 + 13 * r.width) / 1e6
    ax.barh(k, embed, color=COLORS["data"], label="embeddings" if k == 0 else None)
    ax.barh(k, layers_part, left=embed, color=COLORS["dec"],
            label="layers" if k == 0 else None)
    share = embed / (embed + layers_part)
    ax.text(embed + layers_part + 20, k, f"embeddings {share:.0%}", va="center",
            color=COLORS["muted"])
ax.set_yticks(range(len(names)), names)
ax.invert_yaxis()
ax.set_xlabel("millions of parameters")
ax.set_xlim(0, 2100)
ax.legend(loc="upper right", frameon=False)
plt.show()

Output of the code above

3. GPT-2 in PyTorch

A GPT-2 layer has the same two parts as every Transformer layer: self-attention, then a feed-forward network, each wrapped in a residual connection. Two things changed since BERT.

The mask is causal. And the LayerNorm moved: BERT normalised after adding the residual, GPT-2 normalises the input of each part and leaves the residual path alone. The GPT-2 paper describes it in one sentence, and adds a final LayerNorm after the last layer.

Post-norm against pre-norm. In BERT's post-norm layer the LayerNorm sits on the main path, after each addition, so every layer's output is renormalised. In GPT-2's pre-norm layer the LayerNorm sits on the branch, before attention and before the feed-forward network; the main path, the residual stream, runs from the embeddings to the final LayerNorm with nothing on it but additions Post-norm against pre-norm. In BERT's post-norm layer the LayerNorm sits on the main path, after each addition, so every layer's output is renormalised. In GPT-2's pre-norm layer the LayerNorm sits on the branch, before attention and before the feed-forward network; the main path, the residual stream, runs from the embeddings to the final LayerNorm with nothing on it but additions

The difference matters for training deep models. In a post-norm stack, the gradient that reaches the first layer has passed through a LayerNorm in every layer above it.

In a pre-norm stack there is a clean path of additions from the loss to the embeddings, the residual stream, and each layer only adds to it. Deep pre-norm models train without the careful learning-rate warm-up that post-norm ones need [9], and every large decoder since GPT-2 is pre-norm.

Key point

In a pre-norm model every layer reads the residual stream and adds its result back to it. The stream is a running draft of everything the model has worked out so far, which is why section 6 and the probes of section 11 can read it after any layer.

Attention is the one from the first course, with the causal mask. We write it out once by hand, and then use PyTorch’s scaled_dot_product_attention, which computes the same thing in one fused kernel, without ever storing the \(T \times T\) matrix of weights. The check:

def attention_by_hand(q, k, v):
    length = q.shape[-2]
    scores = q @ k.transpose(-2, -1) / math.sqrt(q.shape[-1])
    causal = torch.ones(length, length, dtype=torch.bool, device=q.device).tril()
    scores = scores.masked_fill(~causal, float("-inf"))
    return scores.softmax(dim=-1) @ v


# batch 1, 2 heads, 5 positions, 4 dimensions
q, k, v = torch.randn(3, 1, 2, 5, 4).unbind(0)
fused = F.scaled_dot_product_attention(q, k, v, is_causal=True)
print("largest difference:", (attention_by_hand(q, k, v) - fused).abs().max().item())

largest difference: 5.960464477539063e-08

The attention layer keeps an optional cache, a dictionary where it stores the keys and values it has computed. It does nothing until section 5, where it makes generation fast.

class Attention(nn.Module):
    def __init__(self, config):
        super().__init__()
        d = config["n_embd"]
        self.heads = config["n_head"]
        self.qkv = nn.Linear(d, 3 * d)  # queries, keys and values in one matrix
        self.out = nn.Linear(d, d)

    def split(self, x):
        batch, length, d = x.shape
        return x.view(batch, length, self.heads, d // self.heads).transpose(1, 2)

    def forward(self, x, cache=None):
        q, k, v = self.qkv(x).split(x.shape[-1], dim=-1)
        q, k, v = self.split(q), self.split(k), self.split(v)
        if cache is not None and "k" in cache:  # the keys and values of earlier tokens
            k = torch.cat([cache["k"], k], dim=2)
            v = torch.cat([cache["v"], v], dim=2)
        if cache is not None:
            cache["k"], cache["v"] = k, v
        # one new token may look at everything; several new tokens need the causal mask
        y = F.scaled_dot_product_attention(q, k, v, is_causal=q.shape[2] > 1)
        y = y.transpose(1, 2).reshape(x.shape)
        return self.out(y)

The feed-forward network uses GELU, as BERT did, but in the approximation OpenAI’s code used, with a tanh:

\[\text{GELU}(x) \approx \tfrac{1}{2}\,x\left(1 + \tanh\!\left(\sqrt{2/\pi}\,\left(x + 0.044715\,x^3\right)\right)\right).\]

The exact GELU, \(x\,\Phi(x)\), needs the error function, which was slow on the hardware of 2018; the tanh version differs from it by less than 0.001 everywhere. A transplant has to use the same one the weights were trained with, or the outputs drift. PyTorch has both, and the config says which one a model wants:

x = torch.linspace(-6, 6, 10001)
print("largest gap, tanh GELU against exact:",
      f"{(F.gelu(x, approximate='tanh') - F.gelu(x)).abs().max().item():.1e}")


class Block(nn.Module):
    def __init__(self, config):
        super().__init__()
        d, eps = config["n_embd"], config["layer_norm_epsilon"]
        self.norm1 = nn.LayerNorm(d, eps=eps)
        self.attention = Attention(config)
        self.norm2 = nn.LayerNorm(d, eps=eps)
        self.up = nn.Linear(d, 4 * d)
        self.down = nn.Linear(4 * d, d)
        # OpenAI's "gelu_new" is the tanh approximation
        self.approximate = "tanh" if config["activation_function"] == "gelu_new" else "none"

    def forward(self, x, cache=None):
        x = x + self.attention(self.norm1(x), cache)
        hidden = F.gelu(self.up(self.norm2(x)), approximate=self.approximate)
        return x + self.down(hidden)

largest gap, tanh GELU against exact: 4.7e-04

The model around the layers: token embeddings plus learned position embeddings, one vector per position up to the context length, as in BERT; the layers; the final LayerNorm. Then the output layer, which is not a layer at all.

GPT-2 scores every token of the vocabulary by a dot product of the last hidden state with that token’s input embedding: the same matrix reads tokens in and writes them out.

This weight tying saves \(V \times d\) parameters, 38.6 million in the smallest GPT-2, and works because a good representation of “the next token is Paris” should look like the embedding of Paris [10].

forward can also hand back the residual stream after every layer, keep_states=True; sections 6 and 11 read it.

class GPT(nn.Module):
    def __init__(self, config):
        super().__init__()
        d = config["n_embd"]
        self.tokens = nn.Embedding(config["vocab_size"], d)
        self.positions = nn.Embedding(config["n_positions"], d)
        self.blocks = nn.ModuleList(Block(config) for _ in range(config["n_layer"]))
        self.norm = nn.LayerNorm(d, eps=config["layer_norm_epsilon"])

    def forward(self, ids, caches=None, keep_states=False):
        past = 0
        if caches is not None and "k" in caches[0]:
            past = caches[0]["k"].shape[2]  # tokens already in the cache
        positions = torch.arange(past, past + ids.shape[1], device=ids.device)
        x = self.tokens(ids) + self.positions(positions)
        states = [x]
        for i, block in enumerate(self.blocks):
            x = block(x, caches[i] if caches is not None else None)
            states.append(x)
        logits = self.unembed(x)
        return (logits, states) if keep_states else logits

    def unembed(self, x):
        return self.norm(x) @ self.tokens.weight.T  # tied: the input table, read backwards

4. OpenAI’s weights in our code

OpenAI’s checkpoint names its tensors differently, so we write down which goes where. One more difference hides in the shapes.

OpenAI’s code stored its linear layers as Conv1D, a name left over from the GPT-1 code, which keeps the weight as (inputs, outputs), the transpose of nn.Linear’s (outputs, inputs). So four of the matrices of every layer go in transposed, and load_state_dict refuses anything missing or left over:

GPT2 = "openai-community/gpt2"
tokenizer = AutoTokenizer.from_pretrained(GPT2)
with open(hf_hub_download(GPT2, "config.json")) as f:
    config = json.load(f)


def openai_names(config):
    names = {"tokens.weight": "wte.weight", "positions.weight": "wpe.weight",
             "norm.weight": "ln_f.weight", "norm.bias": "ln_f.bias"}
    parts = {"norm1": "ln_1", "attention.qkv": "attn.c_attn",
             "attention.out": "attn.c_proj", "norm2": "ln_2",
             "up": "mlp.c_fc", "down": "mlp.c_proj"}
    for i in range(config["n_layer"]):
        for ours, theirs in parts.items():
            for kind in ("weight", "bias"):
                names[f"blocks.{i}.{ours}.{kind}"] = f"h.{i}.{theirs}.{kind}"
    return names


CONV1D = ("qkv.weight", "out.weight", "up.weight", "down.weight")
weights = load_file(hf_hub_download(GPT2, "model.safetensors"))
state = {}
for ours, theirs in openai_names(config).items():
    tensor = weights[theirs]
    if ours.endswith(CONV1D):
        tensor = tensor.T
    state[ours] = tensor
gpt2 = GPT(config)
gpt2.load_state_dict(state, strict=True)
gpt2 = gpt2.to(device).eval()
left_over = sorted(set(weights) - set(openai_names(config).values()))
print(f"loaded {sum(p.numel() for p in gpt2.parameters()):,} parameters")
example = left_over[0]
print(f"not loaded: {len(left_over)} tensors, e.g. {example}, "
      f"of shape {tuple(weights[example].shape)}")

loaded 124,439,808 parameters
not loaded: 12 tensors, e.g. h.0.attn.bias, of shape (1, 1, 1024, 1024)

The tensors left over are the stored causal masks from section 2. Now the check against Hugging Face’s GPT2LMHeadModel, the reference implementation: the same text through both, and the scores of all 50,257 tokens at every position compared.

reference = GPT2LMHeadModel.from_pretrained(GPT2).to(device).eval()
CHECK = ("In 1997 Deep Blue beat Garry Kasparov, the world chess champion, "
         "in a six-game match.")
check = tokenizer(CHECK, return_tensors="pt").to(device)
with torch.no_grad():
    theirs = reference(**check).logits
    ours = gpt2(check["input_ids"])
gap = (ours - theirs).abs().max().item()
print(f"largest difference in the scores: {gap:.1e}")
print(f"the scores themselves reach {theirs.abs().max():.0f}")
assert gap < 1e-2
del reference

largest difference in the scores: 6.1e-05
the scores themselves reach 181

Rounding, again: the two implementations add the same numbers in a different order. Our hundred lines are GPT-2.

Teacher forcing makes it easy to see what the model thinks of a text: one pass gives the probability it assigned to every token that actually came next. A sentence from a 2016 news story, which GPT-2 may well have read:

def next_token_probabilities(text):
    ids = tokenizer(text, return_tensors="pt")["input_ids"].to(device)
    with torch.no_grad():
        logits = gpt2(ids)
    probs = logits[0, :-1].float().softmax(dim=-1)
    targets = ids[0, 1:]
    chosen = probs[torch.arange(len(targets)), targets]
    return tokenizer.convert_ids_to_tokens(ids[0]), chosen.cpu()


TEXT = ("AlphaGo, a program developed by Google DeepMind, beat Lee Sedol "
        "four games to one in Seoul.")
tokens, chosen = next_token_probabilities(TEXT)
print(f"loss {-chosen.log().mean():.2f}")
print("the probability of each token, given the ones before it:")
labels = [token.replace("Ġ", "␣") for token in tokens[1:]]
heatmap(chosen[None].numpy(), cols=labels, role="dec", vmin=0, vmax=1, annotate=True,
        cell=0.42)
plt.show()

loss 3.88
the probability of each token, given the ones before it:

Output of the code above

A loss of 3.88 is an average of very different guesses. Where the left context gives the next token away, GPT-2 is close to certain: Mind after Deep, the ol of Sedol once Sed is there, to after four games.

Where the sentence carries news, it is lost: Lee, Sed and Seoul get next to nothing, although its training text runs to the end of 2017, well after the match. Those tokens are the information in the sentence, and the loss is highest exactly there. A language model is, among other things, a meter of surprise.

5. Generating text, and the KV cache

Generation is a loop:

scores at the last positionpick a tokenappend itrun again

The simplest pick is the most probable token, greedy decoding. Here is GPT-2 continuing the most famous opening move order in chess, the Ruy Lopez, written the way chess books and websites write it:

@torch.no_grad()
def generate(model, ids, steps, use_cache=True, temperature=0.0, generator=None):
    caches = [{} for _ in model.blocks] if use_cache else None
    new = ids
    for _ in range(steps):
        logits = model(new if use_cache else ids, caches)[:, -1].float()
        if temperature == 0:
            token = logits.argmax(dim=-1, keepdim=True)
        else:
            probs = (logits / temperature).softmax(dim=-1)
            token = torch.multinomial(probs, 1, generator=generator)
        ids = torch.cat([ids, token], dim=1)
        new = token  # with a cache, only the new token goes in
    return ids


RUY_LOPEZ = "1. e4 e5 2. Nf3 Nc6 3. Bb5"
prompt = tokenizer(RUY_LOPEZ, return_tensors="pt")["input_ids"].to(device)
print(tokenizer.decode(generate(gpt2, prompt, 40)[0]))

1. e4 e5 2. Nf3 Nc6 3. Bb5 Nc6 4. Nf3 Nc6 5. Bd3 Nc6 6. Bd2 Nc6 7. Bd3 Nc6 8. Bd2

GPT-2 knows what a game record looks like, a number, a dot, a letter and a square, and nothing about the game.

Its first move, 3…Nc6, sends a knight to the square where Black’s knight already stands; White answers 4.Nf3, the same mistake, and from then on the two sides shuffle pieces that aren’t there. Greedy decoding makes the loop worse: once a pattern has repeated, repeating it again is the likeliest continuation.

Look at what the loop recomputes. At step \(t\) the model runs over all \(t\) tokens, although the keys and values of the first \(t-1\) were computed at the previous step and haven’t changed: causal attention never lets an earlier position see a later one, so nothing about them can change.

The KV cache keeps every layer’s keys and values, and each step feeds only the new token, whose query attends to the cached keys.

Generation without and with a KV cache. Without it, every step runs the whole sequence through the model again, so step t costs work proportional to t and the steps get slower. With it, every layer stores the keys and values of the tokens seen so far; each step runs only the new token, computes its query, key and value, appends the key and value to the cache and attends over it Generation without and with a KV cache. Without it, every step runs the whole sequence through the model again, so step t costs work proportional to t and the steps get slower. With it, every layer stores the keys and values of the tokens seen so far; each step runs only the new token, computes its query, key and value, appends the key and value to the cache and attends over it

The cache must not change the result. Greedy decoding with and without it, the same prompt:

with_cache = generate(gpt2, prompt, 40, use_cache=True)
without_cache = generate(gpt2, prompt, 40, use_cache=False)
print("same tokens with and without the cache:", torch.equal(with_cache, without_cache))
assert torch.equal(with_cache, without_cache)

same tokens with and without the cache: True

Now the time each new token takes, with and without the cache, as the text grows:

@torch.no_grad()
def time_per_token(use_cache, steps):
    caches = [{} for _ in gpt2.blocks] if use_cache else None
    ids = prompt
    new = prompt
    times = []
    for _ in range(steps):
        if device == "cuda":
            torch.cuda.synchronize()
        start = time.perf_counter()
        logits = gpt2(new if use_cache else ids, caches)[:, -1]
        token = logits.argmax(dim=-1, keepdim=True)
        if device == "cuda":
            torch.cuda.synchronize()
        times.append(time.perf_counter() - start)
        ids = torch.cat([ids, token], dim=1)
        new = token
    return np.array(times) * 1000


STEPS = 60 if SMOKE else 900
time_per_token(True, 5)  # warm up the GPU
slow = time_per_token(False, STEPS)
fast = time_per_token(True, STEPS)
print(f"{STEPS} tokens: {slow.sum() / 1000:.1f} s without the cache, "
      f"{fast.sum() / 1000:.1f} s with it")

fig, ax = plt.subplots(figsize=(6.4, 3.0))
length = prompt.shape[1] + np.arange(STEPS)
ax.plot(length, slow, color=COLORS["line"], label="no cache: the whole text every step")
ax.plot(length, fast, color=COLORS["dec"], label="KV cache: one token every step")
ax.set_xlabel("tokens in the text")
ax.set_ylabel("ms for the next token")
ax.set_ylim(bottom=0)
ax.legend(frameon=False)
plt.show()

900 tokens: 36.1 s without the cache, 4.9 s with it

Output of the code above

Without the cache, each step reruns the whole text, so the time per token grows with the length and the total grows with its square; the staircase is the GPU working in blocks of fixed size, so the cost jumps every so many tokens.

With the cache, token 900 costs what token 20 did. For 900 tokens that is 36.1 seconds against 4.9 on a T4, and the gap widens with every token.

The cache costs memory instead: two vectors of width \(d\) per token in every layer. For GPT-2 that is small; for the models of today it is the main thing that limits how many conversations fit on one GPU, which is what lesson 8 of this course is about. In 16-bit numbers:

def kv_bytes_per_token(layers, d):
    return 2 * layers * d * 2  # keys and values, every layer, 2 bytes per number


for name, layers, d in [("GPT-2", 12, 768), ("GPT-2 XL", 48, 1600), ("GPT-3", 96, 12288)]:
    per_token = kv_bytes_per_token(layers, d)
    print(f"{name:9s} {per_token / 1024:7,.0f} KB per token, "
          f"{per_token * 2048 / 2**30:5.2f} GB for 2,048 tokens")

GPT-2          36 KB per token,  0.07 GB for 2,048 tokens
GPT-2 XL      300 KB per token,  0.59 GB for 2,048 tokens
GPT-3       4,608 KB per token,  9.00 GB for 2,048 tokens

GPT-3’s numbers are from its paper; it has no public weights. A single 2,048-token conversation with it holds several gigabytes of cache, before anyone else is served.

Greedy decoding is a poor way to write prose: it picks the safest token every time, and safe texts repeat, as GPT-2’s chess shows.

Sampling picks the next token at random from the model’s distribution, usually sharpened by a temperature \(\tau\): the scores are divided by \(\tau\) before the softmax, so \(\tau < 1\) makes likely tokens likelier and \(\tau \to 0\) is greedy. The same prompt, three samples at \(\tau = 0.8\):

generator = torch.Generator(device=device).manual_seed(SEED)
START = "The strangest chess game ever played"
start = tokenizer(START, return_tensors="pt")["input_ids"].to(device)
for _ in range(3):
    sample = generate(gpt2, start, 30, temperature=0.8, generator=generator)
    print("·", tokenizer.decode(sample[0]).replace("\n", " "))

· The strangest chess game ever played, the only way to receive the prize?  I have been making chess updates since last November, and when it comes to chess, the most
· The strangest chess game ever played for a mere four years in a single day.  I'm not speaking of the grand tournament, here, about the Russian chess grandmaster —
· The strangest chess game ever played is Chess Guides. It's a game of chess, at least, but it could be more apt to be released on consoles right now.

6. The logit lens: GPT-2 changing its mind

The residual stream after every layer is a vector of the same width as the last one, and nothing stops us from applying the final LayerNorm and the unembedding to it early.

What comes out is the model’s guess for the next token as it stands after that layer. The idea is called the logit lens [11], and it works on GPT-2 because a pre-norm stack only ever adds to the stream, so each layer refines the same running guess.

The last word of each prompt below is where we look: after every layer, what does GPT-2 expect next?

@torch.no_grad()
def logit_lens(text):
    ids = tokenizer(text, return_tensors="pt")["input_ids"].to(device)
    _, states = gpt2(ids, keep_states=True)
    rows = []
    for layer, state in enumerate(states):
        probs = gpt2.unembed(state[0, -1]).float().softmax(dim=-1)
        best = probs.argmax().item()
        rows.append((layer, tokenizer.decode(best), probs[best].item()))
    return rows


PROMPTS = ["In 1997, IBM's Deep Blue defeated the world chess champion Garry",
           "Bobby Fischer won the 1972 World Chess Championship against Boris"]
lens = {text: logit_lens(text) for text in PROMPTS}
fig, axes = plt.subplots(1, 2, figsize=(9.6, 4.4), sharey=True)
for ax, text in zip(axes, PROMPTS):
    rows = lens[text]
    layers = [layer for layer, _, _ in rows]
    top = [p for _, _, p in rows]
    ax.barh(layers, top, color=COLORS["dec"], height=0.7)
    for layer, word, p in rows:
        ax.text(p + 0.02, layer, repr(word), va="center", fontsize=9, color=COLORS["ink"])
    ax.set_xlim(0, 1.45)
    ax.set_title("… " + " ".join(text.split()[-5:]), fontsize=10)
    ax.set_xlabel("probability of the top token")
axes[0].set_ylabel("after layer (0 = embeddings)")
axes[0].set_yticks(range(len(lens[PROMPTS[0]])))
axes[0].invert_yaxis()
plt.tight_layout()
plt.show()

…

For Deep Blue’s opponent, the first seven layers just repeat the last word of the prompt, Garry: early layers still carry mostly their input, and the lens reads it back. From layer 8 Kas takes over, wavers once, and ends at full confidence.

The second prompt asks for Bobby Fischer’s opponent of 1972, Boris Spassky. The middle layers bet on Johnson, the Boris of British politics; the last one settles, with little confidence, on Becker, a tennis player.

Spassky never appears. The smallest GPT-2 doesn’t know him, and the lens shows how it fails: by sliding from one famous Boris to another.

The lens is biased: the final LayerNorm and unembedding were trained to read the last layer, not the middle ones, so early guesses look worse than what the layer knows. The tuned lens fixes this by training a small affine map for every layer[12].

For GPT-2 the plain lens reads well enough to show the main thing: the answer is assembled gradually, and the residual stream is a place where you can watch it happen. Section 11 will read the same stream with a probe of our own.

7. A gigabyte of chess

Lichess, the free and open-source chess server, publishes every rated game played on it, month by month, under CC0: no rights reserved. By November 2026 the database held more than 8.3 billion standard games, with every player’s clock recorded since April 2017[13].

Hugging Face mirrors it as Parquet files by year and month. The month we take, January 2025, is split into 72 Parquet files of about a gigabyte each; one is plenty.

LICHESS = "Lichess/standard-chess-games"
SHARD = "data/year=2025/month=01/train-00000-of-00072.parquet"
path = hf_hub_download(LICHESS, SHARD, repo_type="dataset")
lichess = pd.read_parquet(path, columns=["Event", "WhiteElo", "BlackElo", "movetext"])
if SMOKE:
    lichess = lichess.head(100_000)
print(f"{len(lichess):,} games")
print(lichess["movetext"][0][:240], "…")

1,394,617 games
1. Nc3 { [%clk 0:03:00] } 1... e5 { [%clk 0:03:00] } 2. e4 { [%clk 0:02:59] } 2... Bc5 { [%clk 0:02:59] } 3. Qh5 { [%clk 0:02:57] } 3... Nf6 { [%clk 0:02:58] } 4. Qxe5+ { [%clk 0:02:56] } 4... Be7 { [%clk 0:02:55] } 5. d3 { [%clk 0:02:55] } …

Every move comes with the time left on the player’s clock, { [%clk 0:02:59] }, and Black’s moves repeat the move number as 1.... The games are of every kind, from one-minute bullet to classical:

# "Rated Blitz game", "Blitz swiss https://lichess.org/swiss/…" -> "Blitz"
SPEEDS = r"(UltraBullet|Bullet|Blitz|Rapid|Classical|Correspondence)"
kind = lichess["Event"].str.extract(SPEEDS)[0]
table(kind.value_counts().rename_axis("time control").reset_index(name="games"))
time controlgames
Blitz652,645
Bullet505,274
Rapid212,558
UltraBullet13,939
Classical8,846
Correspondence1,355

Bullet games are decided by the clock as much as by the moves, and a model trained on them learns to blunder in time trouble. We keep blitz, rapid and classical games where both players are rated 1800 or more.

Then we strip the clocks and write each game the way Adam Karvonen did for his chess-playing GPT[14]: a ; to mark the start, move numbers glued to White’s moves, nothing else.

CLOCK = re.compile(r"\{[^}]*\}")  # { [%clk 0:02:59] } and the like
BLACK_NUMBER = re.compile(r"\d+\.\.\. ")  # "1... " before Black's moves
GLYPHS = re.compile(r"[?!]")  # annotations such as ?? and !?
NUMBER = re.compile(r"(\d+)\. ")  # "1. e4" -> "1.e4"
RESULT = re.compile(r" ?(1-0|0-1|1/2-1/2|\*)$")


def clean(movetext):
    text = CLOCK.sub("", movetext)
    text = BLACK_NUMBER.sub("", text)
    text = GLYPHS.sub("", text)
    text = NUMBER.sub(r"\1.", text)
    text = " ".join(text.split())  # one space between moves
    text = RESULT.sub("", text)
    return ";" + text


slow = kind.isin(["Blitz", "Rapid", "Classical"])
strong = (lichess["WhiteElo"] >= 1800) & (lichess["BlackElo"] >= 1800)
games = [clean(movetext) for movetext in lichess["movetext"][slow & strong]]
games = [game for game in games if game.count(" ") >= 10]  # at least five moves each
print(f"{len(games):,} games kept, e.g.\n{games[0][:200]} …")

250,635 games kept, e.g.
;1.e4 c6 2.d4 d5 3.exd5 cxd5 4.Bd3 Nf6 5.h3 Nc6 6.c3 e5 7.dxe5 Nxe5 8.Be2 Bd6 9.Nf3 Nxf3+ 10.Bxf3 Be6 11.O-O Qc7 12.Nd2 Rd8 13.Nb3 h6 14.Nd4 O-O 15.Re1 a6 16.Rxe6 fxe6 17.Nxe6 Qa5 18.Nxf8 Rxf8 19.Be3  …

The whole language of these games fits in a few dozen characters:

VOCAB = sorted(set("".join(games)))
print(len(VOCAB), "characters:", "".join(VOCAB))
stoi = {char: i for i, char in enumerate(VOCAB)}
lengths = np.array([len(game) for game in games])
print(f"a game is {np.median(lengths):.0f} characters long in the median, "
      f"{np.percentile(lengths, 90):.0f} at the 90th percentile")

32 characters:  #+-.0123456789;=BKNOQRabcdefghx
a game is 387 characters long in the median, 663 at the 90th percentile

Thirty-two characters: the files a to h, the ranks 1 to 8, the pieces K Q R B N, x for a capture, + and # for check and mate, O-O for castling, = for promotion, the digits of move numbers, the dot, the space and the ;.

A model that reads this alphabet is never told what a knight is. It sees N followed by a square, again and again, in the places where knights went.

We hold out 4,000 games for testing and pack the rest into one long stream of character ids. A training example is a window of the stream that starts at the beginning of a game.

CONTEXT = 512  # characters the model reads at once: most games fit
rng = np.random.default_rng(SEED)
order = rng.permutation(len(games))
test_games = [games[i] for i in order[:4000]]
train_games = [games[i] for i in order[4000:]]

lookup = np.zeros(128, dtype=np.uint8)
for char, i in stoi.items():
    lookup[ord(char)] = i
stream = lookup[np.frombuffer("".join(train_games).encode("ascii"), dtype=np.uint8)]
stream = torch.from_numpy(stream.astype(np.int64))
starts = np.flatnonzero(stream.numpy() == stoi[";"])
starts = starts[starts + CONTEXT + 1 <= len(stream)]  # every window must fit in the stream
print(f"{len(train_games):,} training games, {len(stream) / 1e6:.0f} million characters; "
      f"{len(test_games):,} held out")

246,635 training games, 102 million characters; 4,000 held out

8. GPT-2 at the chessboard

Chess notation is all over the web, in game reports, forums and PGN files, so GPT-2 has read some. In a typical position there are “of the order of 30 legal moves”, as Claude Shannon estimated in 1950, among the thousands of strings that look like a move. How often does GPT-2 hit one?

Part of the trouble is visible before the model runs. GPT-2’s tokenizer was built for English prose, and it cuts chess notation into pieces that don’t follow the moves:

LINE = "1. e4 e5 2. Nf3 Nc6 3. Bb5 a6 4. Ba4 Nf6 5. O-O Be7"
pieces = [tokenizer.decode([i]) for i in tokenizer(LINE)["input_ids"]]
token_chips([pieces], labels=["GPT-2 tokens"], max_width=640)
plt.show()
print(f"10 moves, {len(pieces)} tokens")

Output of the code above

10 moves, 35 tokens

To test a player, chess or language model, we take positions from the held-out games: a game cut at a random move, the moves so far as the prompt, and the question of what comes next. python-chess replays the moves and judges the answer. The same positions will test our own model later, written in its own format.

def moves_of(game):
    return [word.split(".")[-1] for word in game[1:].split(" ")]  # "12.Nf3" -> "Nf3"


def prompt_for(moves, upto, spaced):
    # the game before move number `upto` (counted from 0), ready for that move
    text = ""
    for i in range(upto):
        if i % 2 == 0:
            text += f"{i // 2 + 1}." + (" " if spaced else "")
        text += moves[i] + " "
    if upto % 2 == 0:  # White to move: the move number comes first
        text += f"{upto // 2 + 1}."
    return text


def test_positions(games, n, seed):
    rng = random.Random(seed)
    positions = []
    for game in games:
        moves = moves_of(game)
        upto = rng.randrange(1, min(len(moves), 70))
        if len(prompt_for(moves, upto, spaced=False)) > CONTEXT - 10:
            continue  # longer than our model's context
        board = chess.Board()
        for move in moves[:upto]:
            board.push_san(move)
        positions.append({"moves": moves, "upto": upto, "board": board,
                          "played": moves[upto]})
        if len(positions) == n:
            break
    return positions


positions = test_positions(test_games, 100 if SMOKE else 1000, SEED)
example = positions[0]
moves, upto = example["moves"], example["upto"]
print("GPT-2 gets:    ", repr(prompt_for(moves, upto, spaced=True).rstrip()))
print("our model gets:", repr(";" + prompt_for(moves, upto, spaced=False)))
print("the player played:", example["played"])

GPT-2 gets:     '1. d4 Nf6 2. c4 e6 3. Nf3 d5 4. Bg5 Bb4+ 5. Nbd2 dxc4 6. Qa4+ Nc6 7. a3 Bxd2+ 8. Nxd2 Qxd4 9. Nxc4 O-O 10. Rd1 Qc5 11. Bxf6'
our model gets: ';1.d4 Nf6 2.c4 e6 3.Nf3 d5 4.Bg5 Bb4+ 5.Nbd2 dxc4 6.Qa4+ Nc6 7.a3 Bxd2+ 8.Nxd2 Qxd4 9.Nxc4 O-O 10.Rd1 Qc5 11.Bxf6 '
the player played: gxf6

A move can go wrong in several ways, and python-chess tells them apart. It may not be a move at all; it may name a move no piece can make, a knight to a square no knight reaches, or a pawn through a piece; it may be a move the piece could make if the king were not in check, or not pinned: pseudo-legal, as chess programmers say; or it may fit two pieces at once.

def judge(board, move):
    try:
        board.parse_san(move)
        return "legal"
    except chess.InvalidMoveError:
        return "not a move"
    except chess.AmbiguousMoveError:
        return "ambiguous"
    except chess.IllegalMoveError:
        pass
    pseudo_legal = {board.san(m).rstrip("+#") for m in board.pseudo_legal_moves}
    if move.rstrip("+#") in pseudo_legal:
        return "leaves the king in check"
    return "no piece can do that"


def score(make_move, positions):
    verdicts, same_as_player = [], 0
    for position in positions:
        move = make_move(position)
        verdicts.append(judge(position["board"], move))
        same_as_player += move.rstrip("+#") == position["played"].rstrip("+#")
    shares = {}
    for verdict, count in Counter(verdicts).items():
        shares[verdict] = count / len(positions)
    return shares, same_as_player / len(positions)


@torch.no_grad()
def gpt2_move(position):
    prompt = prompt_for(position["moves"], position["upto"], spaced=True).rstrip()
    ids = tokenizer(prompt, return_tensors="pt")["input_ids"].to(device)
    continuation = tokenizer.decode(generate(gpt2, ids, 5)[0, ids.shape[1]:])
    words = continuation.split()
    return words[0] if words else ""


gpt2_shares, gpt2_same = score(gpt2_move, positions)
print(f"GPT-2: {gpt2_shares.get('legal', 0.0):.1%} legal moves, "
      f"{gpt2_same:.1%} the same move as the player")
table(pd.DataFrame({"verdict": list(gpt2_shares), "share": list(gpt2_shares.values())}))

GPT-2: 27.8% legal moves, 3.2% the same move as the player
verdictshare
no piece can do that0.68
legal0.28
leaves the king in check0.01
not a move0.02
ambiguous0.01

GPT-2 almost always writes something shaped like a move: only 2% of its answers aren’t moves at all. But 68% are moves that no piece on the board can make, and only 28% are legal. It has learned the notation, not the game.

9. A GPT that has read nothing but chess

Now the same GPT class, with a 32-character vocabulary and a smaller body: 8 layers, 512 wide, 8 heads, a context of 512 characters. That is the size of Karvonen’s smaller chess model, which will come in handy.

We initialise it the way GPT-2 was: weights drawn from a normal distribution with a standard deviation of 0.02, and the two matrices of every layer that write into the residual stream scaled down by \(1/\sqrt{2L}\), so that the stream doesn’t grow with depth at the start.

chess_config = {"vocab_size": len(VOCAB), "n_positions": CONTEXT, "n_embd": 512,
                "n_head": 8, "n_layer": 8, "layer_norm_epsilon": 1e-5,
                "activation_function": "gelu_new"}


def init_like_gpt2(model, layers):
    for module in model.modules():
        if isinstance(module, (nn.Linear, nn.Embedding)):
            nn.init.normal_(module.weight, std=0.02)
        if isinstance(module, nn.Linear):
            nn.init.zeros_(module.bias)
    for block in model.blocks:  # the two matrices that add to the residual stream
        nn.init.normal_(block.attention.out.weight, std=0.02 / math.sqrt(2 * layers))
        nn.init.normal_(block.down.weight, std=0.02 / math.sqrt(2 * layers))


torch.manual_seed(SEED)
chess_gpt = GPT(chess_config)
init_like_gpt2(chess_gpt, chess_config["n_layer"])
chess_gpt = chess_gpt.to(device)
print(f"{sum(p.numel() for p in chess_gpt.parameters()) / 1e6:.1f} million parameters")

25.5 million parameters

The training loop is the first course’s, with the settings GPT-2-style models are trained with: AdamW with \(\beta_2 = 0.95\) and weight decay 0.1 on the matrices only, a short warm-up and a cosine decay of the learning rate, gradients clipped at norm 1. The T4 computes fast in 16-bit floats, so the forward pass runs under autocast, and a GradScaler keeps small gradients from rounding to zero in that format. torch.compile turns the model into fused GPU kernels for training; generation and the tests use the plain model. Every few minutes we stop to test the model on 200 held-out positions, using the move function written for it below.

@torch.no_grad()
def chess_move(model, position, generator=None):
    prompt = ";" + prompt_for(position["moves"], position["upto"], spaced=False)
    ids = torch.tensor([[stoi[char] for char in prompt]], device=device)
    caches = [{} for _ in model.blocks]
    move = ""
    for _ in range(8):  # the longest move, "exd8=Q+", has 7 characters
        logits = model(ids, caches)[0, -1].float()
        if generator is None:
            char = logits.argmax().item()
        else:
            char = torch.multinomial(logits.softmax(dim=-1), 1, generator=generator).item()
        if VOCAB[char] == " ":
            break
        move += VOCAB[char]
        ids = torch.tensor([[char]], device=device)
    return move


def windows(rng, size):
    picked = starts[rng.integers(0, len(starts), size)]
    rows = torch.stack([stream[s:s + CONTEXT + 1] for s in picked]).to(device)
    return rows[:, :-1], rows[:, 1:]  # inputs, and the same characters shifted by one


BATCH, PEAK_LR, WARMUP = 64, 1e-3, 300
STEPS = 30 if SMOKE else 3400
EVAL_EVERY = 10 if SMOKE else 500


def learning_rate(step):
    if step < WARMUP:
        return PEAK_LR * (step + 1) / WARMUP
    progress = (step - WARMUP) / max(1, STEPS - WARMUP)
    cosine = 0.5 * (1 + math.cos(math.pi * progress))
    return PEAK_LR * (0.1 + 0.9 * cosine)  # down to a tenth of the peak


matrices = [p for p in chess_gpt.parameters() if p.dim() == 2]
others = [p for p in chess_gpt.parameters() if p.dim() < 2]
groups = [{"params": matrices, "weight_decay": 0.1},
          {"params": others, "weight_decay": 0.0}]
optimizer = torch.optim.AdamW(groups, lr=PEAK_LR, betas=(0.9, 0.95))
scaler = torch.amp.GradScaler(enabled=device == "cuda")
# the compiler fuses the training passes into fewer GPU kernels; the first steps wait for it
compiled = torch.compile(chess_gpt) if device == "cuda" else chess_gpt
quick_test = positions[:200]

rng = np.random.default_rng(SEED)
curve = []
start = time.time()
for step in range(STEPS):
    for group in optimizer.param_groups:
        group["lr"] = learning_rate(step)
    inputs, targets = windows(rng, BATCH)
    with torch.autocast(device, dtype=torch.float16, enabled=device == "cuda"):
        logits = compiled(inputs)
    loss = F.cross_entropy(logits.float().view(-1, len(VOCAB)), targets.reshape(-1))
    optimizer.zero_grad(set_to_none=True)
    scaler.scale(loss).backward()
    scaler.unscale_(optimizer)
    torch.nn.utils.clip_grad_norm_(chess_gpt.parameters(), 1.0)
    scaler.step(optimizer)
    scaler.update()
    if (step + 1) % EVAL_EVERY == 0 or step + 1 == STEPS:
        chess_gpt.eval()
        shares, _ = score(lambda position: chess_move(chess_gpt, position), quick_test)
        chess_gpt.train()
        minutes = (time.time() - start) / 60
        curve.append({"step": step + 1, "minutes": minutes, "loss": loss.item(),
                      "legal": shares.get("legal", 0.0)})
        print(f"step {step + 1:5d}  {minutes:5.1f} min  loss {loss.item():.3f}  "
              f"legal moves {shares.get('legal', 0.0):.1%}")
chess_gpt = chess_gpt.eval()

step   500    4.3 min  loss 0.789  legal moves 52.5%
step  1000    8.2 min  loss 0.611  legal moves 75.0%
step  1500   12.1 min  loss 0.559  legal moves 84.5%
step  2000   16.0 min  loss 0.530  legal moves 87.5%
step  2500   19.8 min  loss 0.482  legal moves 91.0%
step  3000   23.7 min  loss 0.460  legal moves 91.5%
step  3400   26.9 min  loss 0.447  legal moves 92.0%

Twenty-five minutes, 111 million characters, about one pass over the training games. After four minutes the model already finds a legal move in half the test positions, far more often than GPT-2 ever does; after twenty it levels off around nine in ten. Nobody told it what a knight is.

curve = pd.DataFrame(curve)
fig, (left, right) = plt.subplots(1, 2, figsize=(9.6, 3.0))
left.plot(curve["minutes"], curve["loss"], color=COLORS["dec"], marker="o", markersize=3)
left.set_xlabel("minutes of training")
left.set_ylabel("loss per character")
right.plot(curve["minutes"], curve["legal"], color=COLORS["dec"], marker="o", markersize=3,
           label="our chess GPT")
right.axhline(gpt2_shares.get("legal", 0.0), color=COLORS["line"], linestyle="--",
              label="GPT-2")
right.set_xlabel("minutes of training")
right.set_ylabel("legal moves, 200 positions")
right.set_ylim(0, 1)
right.legend(frameon=False, loc="lower right")
plt.tight_layout()
plt.show()

Output of the code above

10. Does it play chess?

The full test, on the same 1,000 positions GPT-2 got, written in our model’s format:

chess_shares, chess_same = score(lambda p: chess_move(chess_gpt, p), positions)
print(f"our chess GPT: {chess_shares.get('legal', 0.0):.1%} legal moves, "
      f"{chess_same:.1%} the same move as the player")

our chess GPT: 89.5% legal moves, 32.5% the same move as the player

Nine moves in ten are legal. And one move in three is the move the human actually played, in positions with some thirty legal options: the model has learned not only the rules, but how players rated 1800 like to play.

What would more training buy? Karvonen published the models of his paper on the Hugging Face Hub under the MIT license, and his 8-layer model is ours to the letter: 8 layers, 512 wide, 8 heads, the same 32 characters.

It was trained with nanoGPT on 16 million Lichess games, and “took 72 hours to train on one RTX 3090 GPU”. Our class can run it, after one more transplant.

Three things differ from GPT-2’s checkpoint: nanoGPT’s linear layers are nn.Linear, so nothing is transposed; they have no biases, so ours are set to zero, which computes the same function; and nanoGPT uses the exact GELU, which goes into the config.

Trap

The file is a PyTorch checkpoint, which is a pickle, and unpickling can run any code the file contains. weights_only=True makes torch.load refuse everything but tensors and plain values; never load a checkpoint from the internet without it.

path = hf_hub_download("adamkarvonen/chess_llms", "lichess_8layers_ckpt_no_optimizer.pt")
checkpoint = torch.load(path, map_location="cpu", weights_only=True)
args = checkpoint["model_args"]
print(args)
seen = checkpoint["iter_num"] * checkpoint["config"]["batch_size"] * args["block_size"]
print(f"characters seen in training: {seen / 1e9:.0f} billion; "
      f"ours: {STEPS * BATCH * CONTEXT / 1e6:.0f} million")

assert "".join(VOCAB) == " #+-.0123456789;=BKNOQRabcdefghx"  # his alphabet, in his order
karvonen_config = {"vocab_size": args["vocab_size"], "n_positions": args["block_size"],
                   "n_embd": args["n_embd"], "n_head": args["n_head"],
                   "n_layer": args["n_layer"], "layer_norm_epsilon": 1e-5,
                   "activation_function": "gelu"}  # the exact GELU
weights = checkpoint["model"]
karvonen = GPT(karvonen_config)
state = karvonen.state_dict()
zeroed = 0
for ours, theirs in openai_names(karvonen_config).items():
    key = "_orig_mod.transformer." + theirs  # torch.compile added "_orig_mod."
    if key in weights:
        state[ours] = weights[key]
    else:  # trained without biases: a zero bias is the same thing
        state[ours] = torch.zeros_like(state[ours])
        zeroed += 1
karvonen.load_state_dict(state, strict=True)
karvonen = karvonen.to(device).eval()
output_layer = weights["_orig_mod.lm_head.weight"]
tied = torch.equal(output_layer, weights["_orig_mod.transformer.wte.weight"])
print(f"{len(weights)} tensors in the file; {zeroed} biases of ours set to zero; "
      f"output layer tied to the embeddings: {tied}")

{'n_layer': 8, 'n_head': 8, 'n_embd': 512, 'block_size': 1023, 'bias': False, 'vocab_size': 32, 'dropout': 0.0}
characters seen in training: 61 billion; ours: 111 million
52 tensors in the file; 49 biases of ours set to zero; output layer tied to the embeddings: True

The same test for all three:

karvonen_shares, karvonen_same = score(lambda p: chess_move(karvonen, p), positions)
verdicts = ["legal", "leaves the king in check", "no piece can do that", "ambiguous",
            "not a move"]
by_model = {"GPT-2, 124M": gpt2_shares, "ours, 25 minutes": chess_shares,
            "Karvonen's, 72 hours": karvonen_shares}
comparison = pd.DataFrame({"verdict": verdicts})
for name, shares in by_model.items():
    comparison[name] = [shares.get(verdict, 0.0) for verdict in verdicts]
display(table(comparison, digits=3))
print(f"the move the player actually made: GPT-2 {gpt2_same:.1%}, ours {chess_same:.1%}, "
      f"Karvonen's {karvonen_same:.1%}")
verdictGPT-2, 124Mours, 25 minutesKarvonen's, 72 hours
legal0.2780.8950.996
leaves the king in check0.0140.0090.001
no piece can do that0.6820.0880.002
ambiguous0.0050.0080.000
not a move0.0210.0000.001
the move the player actually made: GPT-2 3.2%, ours 32.5%, Karvonen's 52.1%

Karvonen’s model makes a legal move in 99.6% of the positions and the player’s own move in more than half of them. Ours, with the same architecture and 550 times less training, is legal 89.5% of the time, and almost all of its mistakes are moves that no piece can make. What it lacks is the board, not the notation.

The same 25 million parameters, 550 times more training: from nine legal moves in ten to all but a few in a thousand.

A legal move in a position from a human game is one thing; a whole game is harder, because the model’s own moves build the position it has to understand. Each model plays White against an opponent that picks a random legal move.

When a model proposes an illegal move, it gets four more tries, sampled instead of greedy; after five failures it forfeits. A game that outgrows the model’s context is stopped as unfinished, and so is one that reaches 150 half-moves.

def play_against_random(model, seed, max_moves=150):
    rng = random.Random(seed)
    generator = torch.Generator(device=device).manual_seed(seed)
    limit = model.positions.num_embeddings  # the model's context
    board = chess.Board()
    moves, illegal = [], 0
    while not board.is_game_over() and len(moves) < max_moves:
        if board.turn == chess.WHITE:
            if len(prompt_for(moves, len(moves), spaced=False)) > limit - 10:
                return "unfinished", moves, illegal
            position = {"moves": moves, "upto": len(moves)}
            move = None
            for attempt in range(5):
                candidate = chess_move(model, position, generator if attempt > 0 else None)
                if judge(board, candidate) == "legal":
                    move = candidate
                    break
                illegal += 1
            if move is None:
                return "forfeit", moves, illegal
        else:
            move = board.san(rng.choice(list(board.legal_moves)))
        move = board.san(board.parse_san(move))  # the standard spelling, + where it checks
        board.push_san(move)
        moves.append(move)
    if not board.is_game_over():
        return "unfinished", moves, illegal
    return board.result(), moves, illegal


played = []
for name, model in [("ours, 25 minutes", chess_gpt), ("Karvonen's, 72 hours", karvonen)]:
    for seed in range(10 if SMOKE else 100):
        result, moves, illegal = play_against_random(model, seed)
        played.append({"model": name, "result": result, "half-moves": len(moves),
                       "illegal tries": illegal, "game": moves})
played = pd.DataFrame(played)
outcomes = pd.crosstab(played["model"], played["result"])
outcomes["illegal tries per game"] = played.groupby("model")["illegal tries"].mean()
table(outcomes.reset_index())
model1-01/2-1/2forfeitunfinishedillegal tries per game
Karvonen's, 72 hours7418531.00
ours, 25 minutes81494213.95

Here the gap becomes a canyon. Karvonen’s model mates the random opponent in 74 games of 100 and draws 18. Ours forfeits about half of its games, five illegal moves in a row, and most of the rest run out of its context before anyone wins.

The likely reason is the opponent. A random player builds positions that no 1800-rated human would allow, with pieces hanging and kings wandering, and our model has never seen anything like them: its 89.5% on human positions owed a good deal to the habits of human games.

Karvonen’s model has read 61 billion characters, enough to keep its bearings when the game leaves the map. Carlini noticed the same in gpt-3.5-turbo-instruct: given an implausible history that reaches the same position, “it’s only half as good!”

The shortest win of the lot, and where it ended:

# White in capitals, Black in small letters; class 0 is an empty square
PIECES = "PNBRQKpnbrqk"


def board_classes(board):
    classes = np.zeros(64, dtype=np.int64)  # square 0 is a1, 7 is h1, 63 is h8
    for square, piece in board.piece_map().items():
        classes[square] = 1 + PIECES.index(piece.symbol())
    return classes


def draw_board(ax, classes, title, truth=None):
    for square in range(64):
        file, rank = square % 8, square // 8
        light = (file + rank) % 2 == 1
        color = COLORS["card"] if light else COLORS["grid"]
        ax.add_patch(plt.Rectangle((file, rank), 1, 1, color=color))
        if truth is not None and classes[square] != truth[square]:
            ax.add_patch(plt.Rectangle((file + 0.04, rank + 0.04), 0.92, 0.92, fill=False,
                                       edgecolor=COLORS["dec"], linewidth=2))
        if classes[square] > 0:
            piece = PIECES[classes[square] - 1]
            ax.text(file + 0.5, rank + 0.47, piece, ha="center", va="center", fontsize=13,
                    weight="bold", family="monospace",
                    color=COLORS["ink"] if piece.isupper() else COLORS["muted"])
    ax.set_xlim(0, 8)
    ax.set_ylim(0, 8)
    ax.set_xticks(np.arange(8) + 0.5, list("abcdefgh"))
    ax.set_yticks(np.arange(8) + 0.5, range(1, 9))
    ax.tick_params(length=0)
    ax.set_aspect("equal")
    ax.grid(False)
    ax.set_title(title, fontsize=10)


wins = played[played["result"] == "1-0"]
if len(wins):
    shortest = wins.loc[wins["half-moves"].idxmin()]
    game = shortest["game"]
    print(shortest["model"] + ":", prompt_for(game, len(game), spaced=True).strip())
    board = chess.Board()
    for move in game:
        board.push_san(move)
    fig, ax = plt.subplots(figsize=(3.2, 3.4))
    draw_board(ax, board_classes(board), f"after {len(game)} half-moves: checkmate")
    plt.show()

Karvonen's, 72 hours: 1. e4 f5 2. exf5 a6 3. Qh5+ g6 4. fxg6 Nc6 5. g7#

Output of the code above

11. Where is the board?

The model writes legal moves most of the time, and a legal move depends on where every piece stands. Nothing in its input says where they stand: the text lists moves, and a square is mentioned only when a piece goes there.

To get Nc6 right after twenty moves, the model has to have kept track of which knight is where and which squares are free. Does it keep a board somewhere, or has it found some other trick?

The question has a history. In 2022 Kenneth Li and colleagues trained an 8-layer GPT on the moves of Othello, a game with “a vocabulary of 60 words”, and asked whether it kept the board [15].

Linear probes failed: their error rates “never dip below 20%”, and the paper concluded the board was there in a nonlinear form.

A year later Neel Nanda and colleagues found that the board was linear after all, once you ask the right question: the model doesn’t store squares as black and white, but “relative to the current player of each timestep (MINE, YOURS, EMPTY)”, and a linear probe for that reads it at 99.6%[16].

Adam Karvonen then did the same for chess, with a model like ours trained for 72 hours on 16 million Lichess games, and read the board at 99.1% in its 8-layer model[14]. Ours trained for 25 minutes.

A linear probe is the simplest way to ask. Take the residual stream after some layer at a position in the game, a vector \(h \in \mathbb{R}^{512}\), and train a linear map \(W h + b\) to predict, for every one of the 64 squares, which of 13 things stands on it: nothing, or one of the six White or six Black pieces.

If a linear map can read the board off the stream, the board is there, in the plainest form a network can hold anything: as directions in its space[17]. The probe is trained on states of the frozen model; the model’s weights don’t move.

We read the stream at the dot after each move number, where White is about to move, from games the model never trained on.

At that moment “mine” is always White, so the two ways of labelling the board coincide. The same probe trained on an untrained model with the same architecture is the control: whatever it reads comes from the input itself, not from anything the model learned.

@torch.no_grad()
def residual_states(model, games):
    features = [[] for _ in range(len(model.blocks) + 1)]
    labels = []
    for game in games:
        board = chess.Board()
        text, dots, boards = ";", [], []
        for i, move in enumerate(moves_of(game)):
            if i % 2 == 0:
                text += f"{i // 2 + 1}."
                if len(text) > CONTEXT:
                    break
                dots.append(len(text) - 1)  # the dot: White is about to move
                boards.append(board_classes(board))
            text += move + " "
            board.push_san(move)
        ids = torch.tensor([[stoi[char] for char in text[:CONTEXT]]], device=device)
        _, states = model(ids, keep_states=True)
        for layer, state in enumerate(states):
            features[layer].append(state[0, dots].float().cpu())
        labels.append(torch.from_numpy(np.stack(boards)))
    return [torch.cat(layer) for layer in features], torch.cat(labels)


def train_probe(features, labels, epochs=12):
    mean, std = features.mean(dim=0), features.std(dim=0) + 1e-5  # standardised features
    x, y = ((features - mean) / std).to(device), labels.to(device)
    probe = nn.Linear(features.shape[1], 64 * 13).to(device)
    optimizer = torch.optim.AdamW(probe.parameters(), lr=3e-3, weight_decay=0.01)
    for _ in range(epochs):
        for batch in torch.randperm(len(x), device=device).split(1024):
            logits = probe(x[batch]).view(-1, 13)
            loss = F.cross_entropy(logits, y[batch].view(-1))
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
    def read(h):
        logits = probe(((h - mean) / std).to(device)).view(-1, 64, 13)
        return logits.argmax(dim=-1).cpu()
    return read


def probe_scores(read, features, labels):
    with torch.no_grad():
        predicted = read(features)
    correct = predicted == labels
    return {"all squares": correct.float().mean().item(),
            "occupied squares": correct[labels > 0].float().mean().item(),
            "whole board right": correct.all(dim=1).float().mean().item()}


probe_train, probe_test = test_games[:400 if SMOKE else 3000], test_games[3000:3600]
if SMOKE:
    probe_test = test_games[400:480]
torch.manual_seed(SEED)
untrained = GPT(chess_config)
init_like_gpt2(untrained, chess_config["n_layer"])
untrained = untrained.to(device).eval()

MODELS = {"untrained": untrained, "ours, 25 minutes": chess_gpt,
          "Karvonen's, 72 hours": karvonen}
rows, readers, test_features = [], {}, {}
for name, model in MODELS.items():
    train_x, train_y = residual_states(model, probe_train)
    test_x, test_y = residual_states(model, probe_test)
    for layer in range(len(train_x)):
        read = train_probe(train_x[layer], train_y)
        readers[name, layer] = read
        scores = probe_scores(read, test_x[layer], test_y)
        rows.append({"model": name, "layer": layer, **scores})
    test_features[name] = test_x
probes = pd.DataFrame(rows)
print(f"{len(train_y):,} positions to train each probe, {len(test_y):,} to test it")

103,308 positions to train each probe, 20,693 to test it

How much of the board each layer gives away, on the squares that hold a piece; the empty ones are easy, since most of the board is empty most of the time:

fig, ax = plt.subplots(figsize=(6.4, 3.4))
styles = {"untrained": COLORS["line"], "ours, 25 minutes": COLORS["data"],
          "Karvonen's, 72 hours": COLORS["dec"]}
for name, color in styles.items():
    part = probes[probes["model"] == name]
    ax.plot(part["layer"], part["occupied squares"], marker="o", color=color, label=name)
ax.set_xlabel("residual stream after layer (0 = embeddings)")
ax.set_ylabel("probe accuracy, occupied squares")
ax.set_ylim(0, 1)
ax.legend(frameon=False, loc="lower right")
plt.show()
best_rows = probes.loc[probes.groupby("model")["occupied squares"].idxmax()]
table(best_rows.set_index("model").loc[list(MODELS)].reset_index(), digits=3)

Output of the code above

modellayerall squaresoccupied squareswhole board right
untrained70.7690.5230.055
ours, 25 minutes80.8420.6910.113
Karvonen's, 72 hours60.9840.9720.523

The untrained model sets the floor. About half of the occupied squares can be read from any network’s stream, because the input gives a lot away: a piece that has never moved is where it started, and the last move says where one piece went.

Karvonen’s model climbs far above that floor in the middle of its stack. After layer 6 the probe reads 97.2% of the occupied squares, 98.4% of all squares, and the whole board without a single mistake in 52% of the positions; the last two layers give a little back as they turn the board into the next move.

His paper found the peak at the same place, “the 8 layer model obtains peak accuracy at layer 6”.

Ours, after 25 minutes, reads 69% of the occupied squares, and its curve is still rising at the last layer: the board is under construction. Nobody asked either model for a board. Predicting the next character of a game record was enough to make one, and a single matrix can read it off.

What the probes see, on one position from a held-out game: the real board, and the board read off the residual stream of each trained model, at the layer where Karvonen’s model holds it best. Squares a probe gets wrong are outlined.

layer = int(best_rows.set_index("model").loc["Karvonen's, 72 hours", "layer"])
index = len(test_y) // 2
truth = test_y[index].numpy()
fig, axes = plt.subplots(1, 3, figsize=(9.6, 3.6))
draw_board(axes[0], truth, "the real board")
for ax, name in zip(axes[1:], ["ours, 25 minutes", "Karvonen's, 72 hours"]):
    features = test_features[name][layer][index:index + 1]
    board_read = readers[name, layer](features)[0].numpy()
    draw_board(ax, board_read, f"{name}, layer {layer}", truth)
plt.tight_layout()
plt.show()

Output of the code above

Karvonen’s model reads this position without a mistake. Ours gets all of it but one pawn, and the mistake is telling: the pawn went from a2 to a4 long ago, and our probe still finds it on a2, its square at the start of the game.

A model that has half learned to track the board falls back on what boards usually look like.

12. Where GPT is the right tool today

GPT-2 itself has retired from every job but two. It is the laboratory animal of interpretability: open, small enough to run anywhere, and studied so thoroughly that new methods are usually shown on it first, the way the logit lens was.

And it is the reference decoder that people rebuild to learn how one works, as you just did. For writing, answering and following instructions, use a modern open model; lesson 6 is about what changed in them.

The decoder-only design, on the other hand, is everywhere. Every chat model is one, and so are the models that write code or treat proteins, music or clicks as sequences of tokens.

The chess model shows why the design spread: give a GPT a sequence of symbols with structure, and next-token prediction alone makes it learn the structure, here the rules and the board, without anyone writing them down. It also shows the limit.

Predicting the next move of players rated 1800 is not the same as playing well; Carlini’s gpt-3.5-turbo-instruct answered one of the worst openings in chess with the same opening.

It’s important to remember, though, that this model is not playing to win. It’s playing to maximize the likelihood of the PGN text that’s been provided.

Nicholas Carlini, 2023

DeepMind trained a 270-million-parameter Transformer on Stockfish’s evaluations of positions instead of human moves, and it reached a Lichess blitz rating of 2895 against humans without any search[2].

When Kaggle and Google DeepMind ran a chess tournament for frontier language models in August 2025, o3 beat Grok 4 4–0 in the final, and the organisers wrote that engines like Stockfish “would beat every frontier model without a doubt” [18]. A language model plays the game its data plays.

GPT-2’s passport, and the family tree with this lesson’s branch:

GPT-2's passport. Published February 2019 by OpenAI, the full model released in November 2019. Type: decoder only. Pre-training: next-token prediction on WebText, 40 GB of web pages. Attention: causal, 12 heads per layer in the smallest model. Positions: learned, up to 1,024 tokens. Normalisation: pre-norm LayerNorm and a final one. Activation: GELU, tanh approximation. Vocabulary: byte-level BPE, 50,257 tokens, output tied to the input embeddings. Sizes: 124 million, 355 million, 774 million and 1.5 billion parameters. Good at: interpretability research, teaching, small baselines. Not for: facts, instructions, long inputs GPT-2's passport. Published February 2019 by OpenAI, the full model released in November 2019. Type: decoder only. Pre-training: next-token prediction on WebText, 40 GB of web pages. Attention: causal, 12 heads per layer in the smallest model. Positions: learned, up to 1,024 tokens. Normalisation: pre-norm LayerNorm and a final one. Activation: GELU, tanh approximation. Vocabulary: byte-level BPE, 50,257 tokens, output tied to the input embeddings. Sizes: 124 million, 355 million, 774 million and 1.5 billion parameters. Good at: interpretability research, teaching, small baselines. Not for: facts, instructions, long inputs

The family of the Transformer as this course covers it, with lesson 4, GPT, highlighted and lessons 1 to 3 done The family of the Transformer as this course covers it, with lesson 4, GPT, highlighted and lessons 1 to 3 done

References

  1. Carlini, N. (2023). Playing chess with large language models. ↩
  2. Ruoss, A., Delétang, G., Medapati, S., Grau-Moya, J., Wenliang, L. K., Catt, E., Reid, J., Lewis, C. A., Veness, J., & Genewein, T. (2024). Amortized planning with large-scale transformers: a case study on chess. NeurIPS 2024. ↩
  3. Dynomight (2024). Something weird is happening with LLMs and chess. ↩
  4. Burns, C., Izmailov, P., Kirchner, J. H., Baker, B., Gao, L., Aschenbrenner, L., Chen, Y., Ecoffet, A., Joglekar, M., Leike, J., Sutskever, I., & Wu, J. (2023). Weak-to-strong generalization: eliciting strong capabilities with weak supervision. arXiv:2312.09390. ↩
  5. Radford, A., Narasimhan, K., Salimans, T., & Sutskever, I. (2018). Improving language understanding by generative pre-training. OpenAI. ↩
  6. OpenAI, 2019. openai.com. ↩
  7. Radford, A., Wu, J., Child, R., Luan, D., Amodei, D., & Sutskever, I. (2019). Language models are unsupervised multitask learners. OpenAI. ↩
  8. Brown, T. B., Mann, B., Ryder, N., Subbiah, M., et al. (2020). Language models are few-shot learners. NeurIPS 2020. ↩
  9. Xiong, R., Yang, Y., He, D., Zheng, K., Zheng, S., Xing, C., Zhang, H., Lan, Y., Wang, L., & Liu, T.-Y. (2020). On layer normalization in the Transformer architecture. ICML 2020. ↩
  10. Press, O., & Wolf, L. (2017). Using the output embedding to improve language models. EACL 2017. ↩
  11. nostalgebraist (2020). interpreting GPT: the logit lens. LessWrong. ↩
  12. Belrose, N., Ostrovsky, I., McKinney, L., Furman, Z., Smith, L., Halawi, D., Biderman, S., & Steinhardt, J. (2023). Eliciting latent predictions from transformers with the tuned lens. arXiv:2303.08112. ↩
  13. Lichess (2026). Lichess open database. CC0. ↩
  14. Karvonen, A. (2024). Emergent world models and latent variable estimation in chess-playing language models. COLM 2024. ↩
  15. Li, K., Hopkins, A. K., Bau, D., Viégas, F., Pfister, H., & Wattenberg, M. (2023). Emergent world representations: exploring a sequence model trained on a synthetic task. ICLR 2023. ↩
  16. Nanda, N., Lee, A., & Wattenberg, M. (2023). Emergent linear representations in world models of self-supervised sequence models. BlackboxNLP 2023. ↩
  17. Alain, G., & Bengio, Y. (2016). Understanding intermediate layers using linear classifier probes. arXiv:1610.01644. ↩
  18. Google, 2025. blog.google. ↩
  19. Toshniwal, S., Wiseman, S., Livescu, K., & Gimpel, K. (2022). Chess as a testbed for language model state tracking. AAAI 2022.
  20. Shannon, C. E. (1950). Programming a computer for playing chess. Philosophical Magazine, 41(314).
]]> Eva Koroleva Embeddings and search: Sentence-BERT, cross-encoders and ColBERT 2026-04-08T12:00:00+00:00 2026-11-02T00:00:00+00:00 https://blog.korolevalab.com/blog/2026/transformer-architectures-03-search Transformer architectures: from BERT to diffusion · lesson 3 of 15 · all lessons · run on Kaggle · open in Colab · code on GitHub

In 2019 BERT was the best reader of English and one of the worst search engines. Ask it which two of 10,000 sentences mean the most similar things, and it has to read every pair, which “requires about 50 million inference computations (~65 hours)”.

Nils Reimers and Iryna Gurevych at TU Darmstadt changed one thing: they made BERT write down what it read as one vector per sentence, so that comparing two texts became a dot product. The same search then took “about 5 seconds”[1].

Their Sentence-BERT now sits, under various names, behind most of the search you meet: the retrieval half of every RAG system, “related articles”, duplicate detection, the search box of a documentation site. It also has rivals with very different trade-offs, and a thirty-year-old formula that none of them can afford to ignore.

Three ways to score a query against a document. A cross-encoder reads the two texts joined into one input and outputs a relevance score: one full pass of the encoder for every pair. A bi-encoder, Sentence-BERT, encodes each text separately into one vector and compares the vectors by cosine, so the documents can be encoded before any query arrives. Late interaction, ColBERT, also encodes them separately but keeps one vector per token and sums, over the query's tokens, the best match in the document Three ways to score a query against a document. A cross-encoder reads the two texts joined into one input and outputs a relevance score: one full pass of the encoder for every pair. A bi-encoder, Sentence-BERT, encodes each text separately into one vector and compares the vectors by cosine, so the documents can be encoded before any query arrives. Late interaction, ColBERT, also encodes them separately but keeps one vector per token and sums, over the query's tokens, the best match in the document

In this lesson you’ll build all three, check each against the real weights, and pit them against each other on a task where good search is easy to tell from bad: finding machine-learning papers from the questions researchers actually ask.

What you’ll do

  1. Write BM25, the formula every neural retriever still has to beat, and see how far counting words goes.
  2. See why an encoder’s own vectors are useless for search, and fix that the way Sentence-BERT does: train a bi-encoder with in-batch negatives, on nothing but the titles and abstracts of the papers themselves.
  3. Load the most downloaded model on the Hugging Face Hub into the same code, find out what a billion training pairs buy, and adapt it to the corpus in ten minutes.
  4. Write a cross-encoder reranker and ColBERT’s late interaction, and check both against the real weights.
  5. Assemble a search engine from the pieces, and measure what each costs in quality, time and memory.

What you need: a GPU. On Kaggle, Settings → Accelerator → GPU; on Colab, Runtime → Change runtime type → T4 GPU. The notebook downloads five small encoders and the titles and abstracts of 64,000 papers, and trains two encoders for ten minutes each; on a T4 it runs in about 45 minutes.

Setup

One small helper package, tfs, draws the course’s plots and token chips and knows where the notebook runs (Kaggle, Colab or your own machine). Everything the lesson teaches is written out in the notebook itself.

import importlib.util
import subprocess
import sys

TFS = "git+https://github.com/xmarva/transformers-from-scratch@main"
if importlib.util.find_spec("tfs") is None:
    subprocess.check_call([sys.executable, "-m", "pip", "install", "-q", TFS])

from tfs import env
from tfs.viz import table, token_chips, use_course_style

env.quiet()  # no progress bars or environment warnings in the outputs
SEED = 42
env.seed_everything(SEED)
COLORS = use_course_style()
WORK_DIR = env.work_dir()

import json
import math
import re
import string
import time
from collections import Counter, defaultdict

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import torch
import torch.nn as nn
import torch.nn.functional as F
from huggingface_hub import hf_hub_download
from IPython.display import display
from safetensors.torch import load_file
from tfs.viz import bars, heatmap
from transformers import AutoTokenizer
# the reference implementations, for the checks
from transformers import AutoModelForSequenceClassification, BertModel

SMOKE = env.smoke()  # CI runs the lesson on a small slice; you run it in full
device = "cuda" if torch.cuda.is_available() else "cpu"
assert device == "cuda" or SMOKE, "This lesson encodes 58,000 papers six times: use a GPU."
if device == "cuda":
    print("GPU:", torch.cuda.get_device_name(0))
half = torch.float16 if device == "cuda" else torch.float32  # what indexes are stored in

GPU: Tesla T4

1. The task: 597 questions, 64,183 papers

LitSearch is a benchmark built at Princeton from the way researchers look for papers [2]. Some of its 597 questions were written by GPT-4 from paragraphs of real papers that cite other work, and ask for the paper behind the citation; the others were written by authors of ACL 2023 and ICLR 2024 papers about their own papers.

The benchmark’s authors checked every question by hand. The corpus is 64,183 papers from the ACL Anthology and ICLR, and each question names the one or few papers it is after.

The benchmark was not built to flatter anyone: in its paper, Google Search and research tools lagged behind the best neural retriever “by up to 32 recall points”.

The corpus files also hold the full text of every paper, more than a gigabyte of it. We need only titles and abstracts, and Parquet stores every column separately, so we read just those three columns straight from the Hugging Face Hub: about 35 MB instead of 1.2 GB.

LITSEARCH = "hf://datasets/princeton-nlp/LitSearch/"
queries = pd.read_parquet(LITSEARCH + "query/full-00000-of-00001.parquet")
shards = [pd.read_parquet(LITSEARCH + f"corpus_clean/full-{k:05d}-of-00006.parquet",
                          columns=["corpusid", "title", "abstract"])
          for k in range(1 if SMOKE else 6)]
papers = pd.concat(shards, ignore_index=True)
print(f"{len(queries)} questions, {len(papers):,} papers")

has_text = (papers["title"].str.len() > 0) | (papers["abstract"].str.len() > 0)
print(f"{(~has_text).sum():,} papers have neither a title nor an abstract: we drop them")
papers = papers[has_text].reset_index(drop=True)
if SMOKE:  # a small corpus that still holds the answers to a few questions
    wanted = {i for ids in queries["corpusids"] for i in ids}
    keep = papers["corpusid"].isin(wanted) | (papers.index < 600)
    papers = papers[keep].reset_index(drop=True)

titles = papers["title"].str.strip().tolist()
texts = [f"{t}. {a}" if t and a else t or a
         for t, a in zip(titles, papers["abstract"].str.strip())]
row_of = {paper: i for i, paper in enumerate(papers["corpusid"])}
queries["gold"] = [[row_of[i] for i in ids if i in row_of] for ids in queries["corpusids"]]
queries = queries[queries["gold"].map(len) > 0].reset_index(drop=True)
if SMOKE:
    queries = queries.head(24)
questions, gold = queries["query"].tolist(), queries["gold"].tolist()
print(f"{len(questions)} questions keep at least one paper to find, among {len(texts):,}")

597 questions, 64,183 papers
6,197 papers have neither a title nor an abstract: we drop them
593 questions keep at least one paper to find, among 57,986

A few questions with the papers they want:

examples = queries.groupby("query_set").head(1)
table(pd.DataFrame({"question": examples["query"],
                    "the paper it wants": [titles[g[0]] for g in examples["gold"]],
                    "set": examples["query_set"]}))
questionthe paper it wantsset
Are there any research papers on methods to compress large-scale language models using task-agnostic knowledge distillation techniques?TinyBERT: Distilling BERT for Natural Language Understandinginline_acl
*Could you suggest a dataset with legally or ethically contentious content, and labels for acceptable and non-acceptable questions.SQUARE: A Large-Scale Dataset of Sensitive Questions and Acceptable Responses Created Through Human-Machine Collaborationinline_nonacl
Are there any examples of using dense phrase retrieval systems in the automatic curation of entity dictionaries?Automatic Creation of Named Entity Recognition Datasets by Querying Phrase Representationsmanual_acl
Can we reduce visual tokens in vision transformers right from the beginning?SparseFormer: Sparse Visual Recognition via Limited Latent Tokensmanual_iclr

The questions sound like a colleague asking in a corridor. Some are specific, “is there a paper that applies large language models to visual Raven’s Progressive Matrices?”; some are broad enough to have several right answers, and LitSearch marks those with specificity = 0.

The titles in the corpus are as they were extracted from the PDFs, sometimes in capitals and sometimes with a conference header in front, which is what real corpora look like.

LitSearch scores a search engine by recall@k: of the papers a question asks for, the share found in the top k results.

Most questions want exactly one paper, so recall@5 is close to “how often the right paper is on the first screen”. Next to it we’ll record what every method costs: milliseconds per question and the size of its index.

results = []


def evaluate(name, scores, seconds, index_bytes):
    """Recall@5 and @20, from the score of every paper for every question."""
    top = torch.as_tensor(scores).topk(20, dim=1).indices.cpu().numpy()
    row = {"method": name}
    for k in (5, 20):
        found = [len(set(g) & set(t[:k])) / len(g) for g, t in zip(gold, top)]
        row[f"recall@{k}"] = np.mean(found)
    row["ms per question"] = 1000 * seconds / len(questions)
    row["index, MB"] = index_bytes / 1e6
    results.append(row)
    return table(pd.DataFrame([row]))

2. BM25: counting words, carefully

Before any neural network, search engines scored a document by the words it shares with the query, and the formula most of them settled on is BM25, from Stephen Robertson’s Okapi system at City University London in the 1990s; its report for TREC-3 says, with no ceremony, “we in effect combined BM11 and BM15 into a single function BM25”[3].

Thirty years later it is the default ranking of Lucene, Elasticsearch and OpenSearch, and the first baseline in every retrieval paper. For a query \(q\) and a document \(d\):

\[\text{BM25}(q, d) = \sum_{w \in q} \text{IDF}(w) \cdot \frac{f(w, d)\,(k_1 + 1)}{f(w, d) + k_1 \left(1 - b + b\,\frac{|d|}{\text{avgdl}}\right)}, \qquad \text{IDF}(w) = \ln\left(1 + \frac{N - n(w) + 0.5}{n(w) + 0.5}\right)\]

where \(f(w, d)\) is how often the word occurs in the document, \(|d|\) the document’s length in words, avgdl the average length, \(N\) the number of documents and \(n(w)\) how many of them contain the word.

BM25's score of a document for a query is a sum over the query's words of three factors: how rare the word is in the collection, the inverse document frequency; how often it occurs in the document, with returns that diminish, controlled by k1; and the document's length relative to the average, controlled by b, so that long documents don't win by repeating everything BM25's score of a document for a query is a sum over the query's words of three factors: how rare the word is in the collection, the inverse document frequency; how often it occurs in the document, with returns that diminish, controlled by k1; and the document's length relative to the average, controlled by b, so that long documents don't win by repeating everything

All three factors are common sense written down. A rare word says more than a common one. The second occurrence of a word matters, the twentieth adds almost nothing.

A long abstract mentions more words by sheer length, so it must mention them more often to score the same. The implementation is an inverted index: for every word, the documents that contain it and how often. A query then touches only the documents that share at least one of its words.

STOP = set("""a an and are as at be been by can for from has have in into is it its
of on or our such than that the their these they this to using was we were which
with""".split())


def words(text):
    return [w for w in re.findall(r"[a-z0-9]+", text.lower()) if w not in STOP]


class BM25:
    def __init__(self, texts, k1=1.2, b=0.75):
        self.k1 = k1
        postings = defaultdict(list)
        self.lengths = np.zeros(len(texts), dtype=np.float32)
        for doc, text in enumerate(texts):
            counts = Counter(words(text))
            self.lengths[doc] = sum(counts.values())
            for word, count in counts.items():
                postings[word].append((doc, count))
        self.index = {}
        for w, p in postings.items():
            docs, counts = zip(*p)
            self.index[w] = (np.array(docs), np.array(counts, dtype=np.float32))
        n = len(texts)
        self.idf = {w: math.log(1 + (n - len(d) + 0.5) / (len(d) + 0.5))
                    for w, (d, _) in self.index.items()}
        # the length part of the denominator, the same for every query
        self.norm = k1 * (1 - b + b * self.lengths / self.lengths.mean())

    def scores(self, query):
        out = np.zeros(len(self.lengths), dtype=np.float32)
        for w in set(words(query)):
            if w in self.index:
                docs, f = self.index[w]
                out[docs] += self.idf[w] * f * (self.k1 + 1) / (f + self.norm[docs])
        return out


start = time.time()
bm25 = BM25(texts)
postings = sum(len(docs) for docs, _ in bm25.index.values())
print(f"index: {len(bm25.index):,} words, {postings:,} postings, "
      f"built in {time.time() - start:.0f} s")

index: 107,892 words, 4,346,702 postings, built in 9 s

Here is what BM25 thinks matters in the first question. Words like papers and methods occur in thousands of abstracts and weigh little; distillation and agnostic are rare and carry the query.

query = questions[0]
weights = {w: bm25.idf.get(w, 0.0) for w in dict.fromkeys(words(query))}
bars(dict(sorted(weights.items(), key=lambda kv: -kv[1])), fmt="{:.1f}",
     title="IDF of each word of the first question")
plt.show()

fig, ax = plt.subplots(figsize=(6.4, 3.4))
f = np.arange(0, 21)
for ratio, color in [(0.5, COLORS["enc"]), (1.0, COLORS["att"]), (2.0, COLORS["dec"])]:
    norm = 1.2 * (1 - 0.75 + 0.75 * ratio)
    ax.plot(f, f * 2.2 / (f + norm), color=color, marker="o", ms=3,
            label=f"document {ratio:g}× the average length")
ax.set_xlabel("occurrences of the word in the document")
ax.set_ylabel("multiplier of the word's IDF")
ax.set_title("Diminishing returns for repetition, a handicap for length")
ax.legend()
plt.show()

IDF of each word of the first question

IDF of each word of the first question

The multiplier can never exceed \(k_1 + 1 = 2.2\), however often the word repeats, and a document twice the average length needs about twice the occurrences to earn the same score. Now the search itself:

start = time.time()
bm25_scores = torch.from_numpy(np.stack([bm25.scores(q) for q in questions]))
evaluate("BM25", bm25_scores, time.time() - start, postings * 8)  # int32 id + float32 count
methodrecall@5recall@20ms per questionindex, MB
BM250.450.601.0334.77

With no training, no GPU and two parameters, BM25 puts the right paper in the top five for almost half the questions, a recall@5 of 0.45.

That is the number every model below has to beat, and the answer to anyone who says keyword search is dead.

3. An encoder’s own vectors: a crowd in a narrow cone

The obvious way to turn an encoder into a search engine is to average the vectors it computes for the tokens of a text and compare the averages by cosine.

Reimers and Gurevych tried it with BERT and found that it “yields rather bad sentence embeddings, often worse than averaging GloVe embeddings”. On their semantic similarity tests, averaged BERT vectors scored 54.8, BERT’s [CLS] vector 29.2, and plain averaged GloVe word vectors from 2014 scored 61.3. Let’s see why.

We take MiniLM-L6, a small encoder with BERT’s architecture. Microsoft distilled a 12-layer, 384-wide MiniLM from a larger model[4], and Nils Reimers, the same one, cut it to six layers by keeping every second layer.

It has five times fewer parameters than BERT-base, and, as section 5 will show, it is the starting point of the most popular embedding model there is. The BERT of the previous lessons loads it unchanged, once it learns to accept the segment ids that the cross-encoder of section 6 will need.

class Embeddings(nn.Module):
    def __init__(self, config):
        super().__init__()
        d = config["hidden_size"]
        self.words = nn.Embedding(config["vocab_size"], d)
        self.positions = nn.Embedding(config["max_position_embeddings"], d)
        self.segments = nn.Embedding(config["type_vocab_size"], d)
        self.norm = nn.LayerNorm(d, eps=config["layer_norm_eps"])
        self.dropout = nn.Dropout(config["hidden_dropout_prob"])

    def forward(self, ids, segments):
        positions = torch.arange(ids.shape[1], device=ids.device)
        x = self.words(ids) + self.positions(positions) + self.segments(segments)
        return self.dropout(self.norm(x))


class Layer(nn.Module):
    def __init__(self, config):
        super().__init__()
        d, eps = config["hidden_size"], config["layer_norm_eps"]
        self.heads = config["num_attention_heads"]
        self.q, self.k, self.v, self.out = (nn.Linear(d, d) for _ in range(4))
        self.norm1 = nn.LayerNorm(d, eps=eps)
        self.up = nn.Linear(d, config["intermediate_size"])
        self.down = nn.Linear(config["intermediate_size"], d)
        self.norm2 = nn.LayerNorm(d, eps=eps)
        self.dropout = nn.Dropout(config["hidden_dropout_prob"])

    def split(self, x):
        batch, length, d = x.shape
        return x.view(batch, length, self.heads, d // self.heads).transpose(1, 2)

    def forward(self, x, padding):
        q, k, v = self.split(self.q(x)), self.split(self.k(x)), self.split(self.v(x))
        scores = q @ k.transpose(-2, -1) / math.sqrt(q.shape[-1])
        lowest = torch.finfo(scores.dtype).min  # not -inf: a row of -inf turns into NaN
        scores = scores.masked_fill(padding[:, None, None, :], lowest)
        y = (self.dropout(scores.softmax(dim=-1)) @ v).transpose(1, 2).reshape(x.shape)
        x = self.norm1(x + self.dropout(self.out(y)))
        return self.norm2(x + self.dropout(self.down(F.gelu(self.up(x)))))


class Encoder(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.embeddings = Embeddings(config)
        n_layers = config["num_hidden_layers"]
        self.layers = nn.ModuleList(Layer(config) for _ in range(n_layers))
        self.dim = config["hidden_size"]

    def forward(self, ids, padding, segments=None):
        segments = torch.zeros_like(ids) if segments is None else segments
        x = self.embeddings(ids, segments)
        for layer in self.layers:
            x = layer(x, padding)
        return x

Every model in this lesson is a BERT inside, and they store their tensors under the same names, with or without a bert. in front. One loader serves them all; whatever a checkpoint keeps besides the encoder, a pooler, a classifier or ColBERT’s projection, it returns separately.

PARTS = {"q": "attention.self.query", "k": "attention.self.key",
         "v": "attention.self.value", "out": "attention.output.dense",
         "norm1": "attention.output.LayerNorm", "up": "intermediate.dense",
         "down": "output.dense", "norm2": "output.LayerNorm"}


def bert_names(config):
    names = {"embeddings.words.weight": "embeddings.word_embeddings.weight",
             "embeddings.positions.weight": "embeddings.position_embeddings.weight",
             "embeddings.segments.weight": "embeddings.token_type_embeddings.weight",
             "embeddings.norm.weight": "embeddings.LayerNorm.weight",
             "embeddings.norm.bias": "embeddings.LayerNorm.bias"}
    for i in range(config["num_hidden_layers"]):
        for ours, theirs in PARTS.items():
            for p in ("weight", "bias"):
                names[f"layers.{i}.{ours}.{p}"] = f"encoder.layer.{i}.{theirs}.{p}"
    return names


def load_encoder(repo, file="model.safetensors"):
    with open(hf_hub_download(repo, "config.json")) as f:
        config = json.load(f)
    path = hf_hub_download(repo, file)
    if file.endswith(".safetensors"):
        weights = load_file(path)
    else:
        weights = torch.load(path, map_location="cpu", weights_only=True)
    prefix = "bert." if "bert.embeddings.word_embeddings.weight" in weights else ""
    encoder = Encoder(config)
    names = bert_names(config)
    encoder.load_state_dict({ours: weights[prefix + t] for ours, t in names.items()})
    used = {prefix + theirs for theirs in names.values()}
    extra = {k.removeprefix(prefix): v for k, v in weights.items()
             if k not in used and not k.endswith("position_ids")}
    return encoder.to(device).eval(), extra


BASE = "nreimers/MiniLM-L6-H384-uncased"
base, _ = load_encoder(BASE, "pytorch_model.bin")
tokenizer = AutoTokenizer.from_pretrained(BASE)
print(f"{sum(p.numel() for p in base.parameters()) / 1e6:.1f}M parameters")

22.6M parameters

An embedding of a text is one vector. We get it by mean pooling, averaging the token vectors while leaving out the padding, and normalise it to length one, so that the dot product of two embeddings is their cosine.

The papers are sorted by length before they are batched: a batch then pads to about the same length, and encoding runs about twice as fast.

def pool(hidden, padding, how):
    if how == "cls":
        return hidden[:, 0]
    keep = (~padding).unsqueeze(-1).to(hidden.dtype)
    return (hidden * keep).sum(1) / keep.sum(1)


@torch.no_grad()
def embed(encoder, texts, max_length, how="mean", prefix="", batch_size=128, tok=None):
    tok = tok or tokenizer
    order = np.argsort([-len(t) for t in texts])
    out = torch.empty(len(texts), encoder.dim, device=device)
    for start in range(0, len(texts), batch_size):
        rows = order[start:start + batch_size]
        batch = tok([prefix + texts[i] for i in rows], padding=True, truncation=True,
                    max_length=max_length, return_tensors="pt").to(device)
        padding = batch["attention_mask"] == 0
        with torch.autocast(device, dtype=torch.float16, enabled=device == "cuda"):
            hidden = encoder(batch["input_ids"], padding)
        vectors = pool(hidden.float(), padding, how)
        out[torch.as_tensor(rows, device=device)] = F.normalize(vectors, dim=-1)
    return out


def dense_search(name, encoder, how="mean", prefix="", tok=None):
    """Encode the corpus once, then every question, and score all papers by cosine."""
    docs = embed(encoder, texts, 256, how, tok=tok)
    start = time.time()
    found = embed(encoder, questions, 64, how, prefix, tok=tok) @ docs.T
    if device == "cuda":
        torch.cuda.synchronize()
    display(evaluate(name, found, time.time() - start, docs.numel() * 2))  # as float16
    return docs, found


start = time.time()
base_docs, base_scores = dense_search("MiniLM-L6, as pre-trained", base)
print(f"{len(texts):,} papers encoded in {time.time() - start:.0f} s")
methodrecall@5recall@20ms per questionindex, MB
MiniLM-L6, as pre-trained0.000.000.2944.53
57,986 papers encoded in 105 s

Not one of the 593 questions finds its paper, not even in the top twenty. The histogram shows why: the cosines of random pairs of papers next to the cosines of questions with the papers they want.

def cosines(docs, found, n=4000):
    rng = torch.Generator().manual_seed(SEED)
    pairs = torch.randint(0, len(docs), (n, 2), generator=rng).to(device)
    random_pairs = (docs[pairs[:, 0]] * docs[pairs[:, 1]]).sum(1).cpu()
    right = torch.stack([found[i, wanted[0]] for i, wanted in enumerate(gold)]).cpu()
    return random_pairs, right


def cosine_histogram(ax, docs, found, title):
    random_pairs, right = cosines(docs, found)
    bins = np.linspace(-0.2, 1.0, 61)
    ax.hist(random_pairs, bins=bins, density=True, color=COLORS["line"], alpha=0.9,
            label="two random papers")
    ax.hist(right, bins=bins, density=True, color=COLORS["enc"], alpha=0.75,
            label="a question and its paper")
    ax.set_title(title)
    ax.set_xlabel("cosine")
    ax.set_yticks([])


fig, ax = plt.subplots(figsize=(6.4, 3.2))
cosine_histogram(ax, base_docs, base_scores, "MiniLM-L6 as pre-trained: all alike")
ax.legend(loc="upper left")
plt.show()

Output of the code above

Two random abstracts have a cosine of about 0.95, and a question with the paper it asks for only about 0.85.

By this model’s measure a question is further from its own answer than any abstract is from any other abstract: the vectors say “this is an abstract” or “this is a question” much louder than what the text is about.

All of them crowd together. Ethayarajh (2019) found the same in almost every layer of BERT, ELMo and GPT-2: their representations are anisotropic, “occupying a narrow cone in the vector space” instead of spreading in all directions.

So what does it return? Here are the papers it puts first most often:

top5 = base_scores.topk(5, dim=1).indices
n_papers = top5.unique().numel()
print(f"{n_papers} different papers fill the top five of {len(questions)} questions")
for paper, count in Counter(top5.flatten().tolist()).most_common(3):
    name = re.sub(r"\s+", " ", titles[paper] or "(no title) " + texts[paper])
    print(f"{count:>4} questions get: {name[:72]}")

556 different papers fill the top five of 593 questions
  82 questions get: Why Neural Translations are the Right Length
  79 questions get: (no title) We detail methods for entity span identification and entity c
  70 questions get: Can language models capture syntactic associations without surface cues?

Most top fives differ, but a few papers turn up everywhere: Why Neural Translations are the Right Length is among the five best answers to 82 of the 593 questions, whatever they ask about. Such points sit near the middle of the crowd and are close to everything; retrieval people call them hubs.

Key point

Nothing in masked language modelling asked the model to put texts with similar meanings close together and the rest far apart; its vectors are good for the next layer, not for being compared with each other. To compare, the encoder has to be trained to compare.

4. Sentence-BERT: teaching an encoder to compare

Sentence-BERT’s architecture is a single encoder applied to each text separately, a bi-encoder (the paper calls it siamese), with mean pooling on top. All the novelty is in the training. In 2019 it learned from pairs of sentences labelled by people as entailment, contradiction or neutral, the SNLI and MultiNLI datasets.

Today the standard is contrastive training with in-batch negatives, which needs no labels at all, only pairs of texts that belong together: a question and its answer, a title and its abstract.

Google’s Smart Reply team put it in one sentence: “we use the responses of other examples in a training batch … as negative responses” [5], and Dense Passage Retrieval made the trick standard for search[6].

Training with in-batch negatives. A batch holds the titles and abstracts of four papers. Every title is compared with every abstract; for each title, its own abstract is the right answer and the other abstracts of the batch are wrong ones, so the loss is cross-entropy over each row of the similarity matrix Training with in-batch negatives. A batch holds the titles and abstracts of four papers. Every title is compared with every abstract; for each title, its own abstract is the right answer and the other abstracts of the batch are wrong ones, so the loss is cross-entropy over each row of the similarity matrix

A batch holds \(B\) pairs \((t_i, a_i)\). Every title is compared with every abstract, and the loss asks each title to pick its own abstract out of the \(B\), as a classification with \(B\) classes:

\[\mathcal{L} = -\frac{1}{B} \sum_{i=1}^{B} \log \frac{\exp\left(s \cdot \cos(t_i, a_i)\right)}{\sum_{j=1}^{B} \exp\left(s \cdot \cos(t_i, a_j)\right)}\]

This loss has many names: InfoNCE[7], the multiple negatives ranking loss, NT-Xent. The scale \(s = 20\), a temperature of 0.05, is what the sentence-transformers library uses by default: cosines live in \([-1, 1]\), and without the scale the softmax over them would be almost flat.

Our pairs are the titles and abstracts of the corpus itself, and that is the practical point. We use no question of LitSearch for training, only the papers we are going to search.

In practice

A company that wants search over its own documents rarely has thousands of labelled queries, but it always has documents, and documents have structure: titles and bodies, questions and answers in a support system, a function’s docstring and its code.

both = [i for i, (t, a) in enumerate(zip(titles, papers["abstract"])) if t and a.strip()]
pairs = [(titles[i], papers["abstract"][i]) for i in both]
BATCH, EPOCHS = (16, 1) if SMOKE else (128, 2)
steps = EPOCHS * (len(pairs) // BATCH)
if SMOKE:
    steps = 6
print(f"{len(pairs):,} title–abstract pairs, {steps:,} steps of {BATCH}")

student, _ = load_encoder(BASE, "pytorch_model.bin")


def pooled(model, texts, max_length):
    batch = tokenizer(texts, padding=True, truncation=True, max_length=max_length,
                      return_tensors="pt").to(device)
    padding = batch["attention_mask"] == 0
    with torch.autocast(device, dtype=torch.float16, enabled=device == "cuda"):
        hidden = model(batch["input_ids"], padding)
    return F.normalize(pool(hidden.float(), padding, "mean"), dim=-1)


def train_bi_encoder(model, scale=20.0, lr=5e-5, warmup=100):
    model.train()
    optimizer = torch.optim.AdamW(model.parameters(), lr=lr)
    schedule = torch.optim.lr_scheduler.LambdaLR(
        optimizer, lambda s: min(1.0, (s + 1) / warmup) * max(0.0, 1 - s / steps))
    scaler = torch.amp.GradScaler(enabled=device == "cuda")
    g = torch.Generator().manual_seed(SEED)
    order = torch.cat([torch.randperm(len(pairs), generator=g) for _ in range(EPOCHS)])
    log = []
    for step in range(steps):
        batch = [pairs[i] for i in order[step * BATCH:(step + 1) * BATCH]]
        t = pooled(model, [title for title, _ in batch], 32)
        a = pooled(model, [abstract for _, abstract in batch], 256)
        logits = scale * t @ a.T  # row i: title i against every abstract of the batch
        target = torch.arange(len(batch), device=device)
        loss = F.cross_entropy(logits, target)
        optimizer.zero_grad()
        scaler.scale(loss).backward()
        scaler.unscale_(optimizer)
        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
        scaler.step(optimizer)
        scaler.update()
        schedule.step()
        accuracy = (logits.argmax(1) == target).float().mean().item()
        log.append((loss.item(), accuracy))
    return model.eval(), np.array(log)


start = time.time()
student, log = train_bi_encoder(student)
print(f"trained in {time.time() - start:.0f} s")

fig, (left, right) = plt.subplots(1, 2, figsize=(8.4, 3.0))
window = max(1, len(log) // 40)
smooth = np.array([log[max(0, i - window):i + 1].mean(0) for i in range(len(log))])
left.plot(smooth[:, 0], color=COLORS["enc"])
left.set_title("loss")
right.plot(smooth[:, 1], color=COLORS["att"])
right.set_title(f"title finds its abstract among {BATCH}")
for ax in (left, right):
    ax.set_xlabel("step")
plt.show()

55,631 title–abstract pairs, 868 steps of 128
trained in 529 s

loss

Within a few hundred steps a title finds its own abstract among 128 nine times out of ten. That is the training task; whether it carries over to questions the model has never seen is what the evaluation says.

student_docs, student_scores = dense_search("MiniLM-L6, trained here", student)
methodrecall@5recall@20ms per questionindex, MB
MiniLM-L6, trained here0.320.460.2544.53

From nothing to a third of the questions in ten minutes, without a single question in the training data. It is still well below BM25, and the reasons are plain: titles are not questions, and 55,000 pairs are not many.

Here is the similarity matrix of eight papers the model trained on, as it looks after training. Six of the eight titles find their own abstract, on the diagonal. The misses are telling: SPREADING VECTORS FOR SIMILARITY SEARCH shares hardly a word with its abstract, which is about quantizers for indexing.

And the bright cells off the diagonal are not noise: the two papers on federated learning, rows 3 and 5, come out close to each other’s abstracts, which is what a search engine should want.

sample = both[:8]
short = [re.sub(r"\s+", " ", titles[i])[:38] for i in sample]
with torch.no_grad():
    t = pooled(student, [titles[i] for i in sample], 32)
    a = pooled(student, [papers["abstract"][i] for i in sample], 256)
heatmap((t @ a.T).cpu().numpy(), rows=short, cols=[f"abstract {k + 1}" for k in range(8)],
        annotate=True, fmt="{:.2f}", cell=0.62, vmin=-0.2, vmax=1.0, role="enc",
        title="Titles against abstracts after training")
plt.show()

fig, axes = plt.subplots(1, 2, figsize=(8.8, 3.0), sharex=True)
cosine_histogram(axes[0], base_docs, base_scores, "before")
cosine_histogram(axes[1], student_docs, student_scores, "after training on our pairs")
axes[0].legend(loc="upper left", fontsize=8)
plt.show()

Titles against abstracts after training

Titles against abstracts after training

The cone has opened. Random pairs of papers now spread around zero, and the questions, which the model never saw during training, sit to the right of them, around 0.5. The model learned what makes two texts in this corpus mean the same thing, from the corpus alone.

5. What a billion pairs buy

all-MiniLM-L6-v2 is not just the most popular embedding model: with 184 million downloads in the month before this lesson was written, it is the most downloaded model of any kind on the Hugging Face Hub.

Its model card says it started from exactly our nreimers/MiniLM-L6-H384-uncased and was trained with the same loss, on 1,170,060,424 pairs from 290 sources: Reddit, Stack Exchange, WikiAnswers, PAQ, and 41,769,185 titles with their abstracts from the S2ORC corpus of scientific papers, which is our task at 750 times our scale. It trained for 100,000 steps of 1,024 pairs on a TPU v3-8.

Same architecture, same starting point, same loss; only the data differ. So the comparison with our model shows what the data buy. Its weights load into our code, and the check against Hugging Face’s BertModel:

MINILM = "sentence-transformers/all-MiniLM-L6-v2"
minilm, _ = load_encoder(MINILM)
reference = BertModel.from_pretrained(MINILM).to(device).eval()
check = tokenizer(questions[:4], padding=True, return_tensors="pt").to(device)
padding = check["attention_mask"] == 0
with torch.no_grad():
    theirs = reference(**check).last_hidden_state
    ours = minilm(check["input_ids"], padding)
gap = (ours - theirs)[~padding].abs().max().item()
print(f"largest difference from Hugging Face: {gap:.1e}")
assert gap < 1e-3
del reference

minilm_docs, minilm_scores = dense_search("all-MiniLM-L6-v2", minilm)

largest difference from Hugging Face: 1.8e-06
methodrecall@5recall@20ms per questionindex, MB
all-MiniLM-L6-v20.420.590.2544.53

Ten points above our ten-minute model, and still three below BM25: 750 times more of our task, and a billion other pairs, bought ten points. Before going further, one more thing that costs nothing to check. The model was trained with mean pooling; here it is read with the [CLS] vector instead, the other common choice:

_ = dense_search("all-MiniLM-L6-v2, [CLS] instead of the mean", minilm, "cls")
methodrecall@5recall@20ms per questionindex, MB
all-MiniLM-L6-v2, [CLS] instead of the mean0.300.430.2544.53

Trap

Twelve points gone, without an error or a warning: the vectors look like vectors, and the search still works, only worse. A model is its weights and its pooling, and the card says which pooling.

The most practical question of the section is the next one. We have a general model trained on a billion pairs, and a corpus with 55,000 pairs of its own. Does a few minutes of the same training on our titles and abstracts, starting from all-MiniLM-L6-v2 instead of from scratch, adapt it to this corpus?

adapted, _ = load_encoder(MINILM)
start = time.time()
adapted, adapted_log = train_bi_encoder(adapted, lr=2e-5)
found = adapted_log[-50:, 1].mean()
print(f"adapted in {time.time() - start:.0f} s; in the last 50 steps a title found "
      f"its abstract {found:.0%} of the time")
adapted_docs, adapted_scores = dense_search("all-MiniLM-L6-v2, adapted here", adapted)

adapted in 527 s; in the last 50 steps a title found its abstract 97% of the time
methodrecall@5recall@20ms per questionindex, MB
all-MiniLM-L6-v2, adapted here0.510.650.2544.53

Ten minutes of the same training took all-MiniLM-L6-v2 from 0.42 to 0.51, above BM25.

The general model brought what a billion pairs teach about language; our pairs showed it what this corpus’s titles and abstracts look like. This is the cheapest improvement in the lesson, and it is open to anyone with a corpus: no labels, no questions, one T4.

BGE models from the Beijing Academy of Artificial Intelligence are a newer generation of the same idea [8]: BERT’s architecture again, with pre-training aimed at retrieval, contrastive training on a larger and cleaner collection, and a final round on labelled data.

The base size has 12 layers and 768 dimensions. Its model card says two things that are easy to miss: the embedding is the vector of [CLS], not the mean, and a query, though not a document, should start with an instruction; without it, the card promises “a slight degradation”.

BGE = "BAAI/bge-base-en-v1.5"
INSTRUCTION = "Represent this sentence for searching relevant passages: "
bge, _ = load_encoder(BGE)
bge_tokenizer = AutoTokenizer.from_pretrained(BGE)
bge_docs, bge_scores = dense_search("bge-base-en-v1.5", bge, "cls", INSTRUCTION,
                                   tok=bge_tokenizer)
_ = dense_search("bge-base-en-v1.5, mean pooling, no instruction", bge,
                 tok=bge_tokenizer)
methodrecall@5recall@20ms per questionindex, MB
bge-base-en-v1.50.500.670.8389.07
methodrecall@5recall@20ms per questionindex, MB
bge-base-en-v1.5, mean pooling, no instruction0.510.660.7189.07

bge-base-en-v1.5, with twice the layers and twice the width, scores 0.50: the same as the small model adapted to our corpus, with an index twice the size and three times the time per question. The second line holds a surprise.

Used against its card’s advice, with mean pooling and no instruction, bge does no worse. Its card warns only of “a slight degradation” without the instruction, and here even that is lost in the noise, while all-MiniLM-L6-v2 read with [CLS] lost twelve points. Pooling matters for one model and not for another, and you can’t know which without checking.

table(pd.DataFrame(results))
methodrecall@5recall@20ms per questionindex, MB
BM250.450.601.0334.77
MiniLM-L6, as pre-trained0.000.000.2944.53
MiniLM-L6, trained here0.320.460.2544.53
all-MiniLM-L6-v20.420.590.2544.53
all-MiniLM-L6-v2, [CLS] instead of the mean0.300.430.2544.53
all-MiniLM-L6-v2, adapted here0.510.650.2544.53
bge-base-en-v1.50.500.670.8389.07
bge-base-en-v1.5, mean pooling, no instruction0.510.660.7189.07

6. The cross-encoder: reading the question and the paper together

A bi-encoder never sees the query and the document at the same time; it compares two summaries written independently.

A cross-encoder is plain BERT reading a pair of texts, in the input format of lesson 1: the query and the document go in as one sequence, [CLS] query [SEP] document [SEP], the two parts marked by their segment ids, and every token of the query attends to every token of the document in every layer.

On top sits a classifier that turns the [CLS] vector into one number, a relevance score.

Putting BERT on top of BM25 this way was one of the first places where BERT took over search [9]. It is the most accurate of the three architectures and the most expensive: nothing can be computed before the query arrives.

Scoring all 58,000 papers for one question would take 58,000 passes of the encoder. So a cross-encoder is used as a reranker: a fast retriever proposes a few dozen candidates and the cross-encoder reorders them.

We take cross-encoder/ms-marco-MiniLM-L6-v2, the same six-layer MiniLM trained on MS MARCO: real Bing queries, each with the passages that human editors marked as containing the answer. Besides the encoder it has BERT’s pooler, a dense layer with tanh on the [CLS] vector, and the classifier with a single output.

class CrossEncoder(nn.Module):
    def __init__(self, repo):
        super().__init__()
        self.encoder, extra = load_encoder(repo)
        d = self.encoder.dim
        self.pooler, self.classifier = nn.Linear(d, d), nn.Linear(d, 1)
        self.pooler.load_state_dict({"weight": extra["pooler.dense.weight"],
                                     "bias": extra["pooler.dense.bias"]})
        self.classifier.load_state_dict({"weight": extra["classifier.weight"],
                                         "bias": extra["classifier.bias"]})

    def forward(self, ids, padding, segments):
        hidden = self.encoder(ids, padding, segments)
        return self.classifier(torch.tanh(self.pooler(hidden[:, 0]))).squeeze(-1)


RERANKER = "cross-encoder/ms-marco-MiniLM-L6-v2"
reranker = CrossEncoder(RERANKER).to(device).eval()
reranker_tokenizer = AutoTokenizer.from_pretrained(RERANKER)
# the reference reads the PyTorch file: this model's safetensors file keeps its tensors at
# offsets that aren't multiples of four bytes; Apple's CPU math library returns NaN there
reference = AutoModelForSequenceClassification.from_pretrained(RERANKER,
                                                               use_safetensors=False)
reference = reference.to(device).eval()
pair = reranker_tokenizer([questions[0]] * 2, [texts[gold[0][0]], texts[0]], padding=True,
                          truncation=True, max_length=256, return_tensors="pt").to(device)
with torch.no_grad():
    theirs = reference(**pair).logits.squeeze(-1)
    ours = reranker(pair["input_ids"], pair["attention_mask"] == 0, pair["token_type_ids"])
print("Hugging Face:", theirs.cpu().numpy().round(3), " ours:", ours.cpu().numpy().round(3))
assert (ours - theirs).abs().max() < 1e-3
del reference
print(reranker_tokenizer.convert_ids_to_tokens(pair["input_ids"][0])[:24])

Hugging Face: [ 1.074 -9.094]  ours: [ 1.074 -9.094]
['[CLS]', 'are', 'there', 'any', 'research', 'papers', 'on', 'methods', 'to', 'com', '##press', 'large', '-', 'scale', 'language', 'models', 'using', 'task', '-', 'ag', '##nostic', 'knowledge', 'di', '##sti']

The first pair is the first question with the paper it wants, the second the same question with a random paper: the scores are logits, and only their order matters. The reranker takes the top 50 papers of a first-stage search and replaces their scores with its own; everything below the 50 stays below.

@torch.no_grad()
def rerank(name, scores, depth=50, max_length=320):
    start = time.time()
    candidates = scores.topk(depth, dim=1).indices
    reranked = torch.full_like(scores, -1e4)
    for i, question in enumerate(questions):
        docs = candidates[i].tolist()
        batch = reranker_tokenizer([question] * depth, [texts[j] for j in docs],
                                   padding=True, truncation="only_second",
                                   max_length=max_length, return_tensors="pt").to(device)
        with torch.autocast(device, dtype=torch.float16, enabled=device == "cuda"):
            logits = reranker(batch["input_ids"], batch["attention_mask"] == 0,
                              batch["token_type_ids"])
        reranked[i, candidates[i]] = logits.float()
    seconds = time.time() - start
    display(evaluate(name, reranked, seconds, 0))
    return reranked


reranked_minilm = rerank("all-MiniLM-L6-v2 → reranker", minilm_scores.to(device))
methodrecall@5recall@20ms per questionindex, MB
all-MiniLM-L6-v2 → reranker0.540.64131.450.00

The reranker lifts all-MiniLM-L6-v2 from 0.42 to 0.54: reading each of the 50 candidates together with the question, it puts the right paper in the first five for about 70 more questions.

It also takes 131 ms per question, five hundred times as long as the search that found the candidates. That is the trade a reranker always offers.

7. ColBERT: a vector for every token

Between the two extremes sits late interaction, from Omar Khattab and Matei Zaharia at Stanford [10]. Like a bi-encoder, ColBERT encodes the query and the document separately, so the documents can be indexed in advance.

Unlike it, ColBERT keeps a vector for every token, shrunk by a linear layer from 768 dimensions to 128 and normalised.

The query and the document meet only at the very end, in MaxSim: every query token finds the document token most similar to it, and the document’s score is the sum of these best matches. Against a BERT reranker, the paper reported search “two orders-of-magnitude faster and requiring four orders-of-magnitude fewer FLOPs per query”.

\[\text{score}(q, d) = \sum_{i \in q} \max_{j \in d} \cos(q_i, d_j)\]

ColBERT's MaxSim. Every query token is compared by cosine with every token of the document; each query token keeps only its best match, and the document's score is the sum of these maxima. Here the query token 'rotary' finds 'rotation' and 'position' finds 'position'; the [MASK] tokens that pad the query also look for a match ColBERT's MaxSim. Every query token is compared by cosine with every token of the document; each query token keeps only its best match, and the document's score is the sum of these maxima. Here the query token 'rotary' finds 'rotation' and 'position' finds 'position'; the [MASK] tokens that pad the query also look for a match

A few details matter for loading the real model, ColBERTv2 [11]. The encoder is BERT-base. The query starts with a marker token after [CLS], and the document with another; ColBERT’s code writes them into BERT’s unused vocabulary slots, [unused0] and [unused1].

A query is always 32 tokens long: what the question doesn’t fill is padded with [MASK], the model doesn’t attend to these masks, but their output vectors still take part in MaxSim.

The paper calls this query augmentation: the masks become soft extra query terms that the model fills with whatever the question implies. Document tokens that are punctuation are dropped from the index.

class ColBERT(nn.Module):
    def __init__(self, repo):
        super().__init__()
        self.encoder, extra = load_encoder(repo)
        self.linear = nn.Linear(self.encoder.dim, 128, bias=False)
        self.linear.load_state_dict({"weight": extra["linear.weight"]})

    def forward(self, ids, padding):
        return F.normalize(self.linear(self.encoder(ids, padding)), dim=-1)


COLBERT = "colbert-ir/colbertv2.0"
colbert = ColBERT(COLBERT).to(device).eval()
colbert_tokenizer = AutoTokenizer.from_pretrained(COLBERT)
Q_MARKER, D_MARKER = colbert_tokenizer.convert_tokens_to_ids(["[unused0]", "[unused1]"])
punctuation = colbert_tokenizer.convert_tokens_to_ids(list(string.punctuation))
PUNCTUATION = torch.tensor(punctuation, device=device)

reference = BertModel.from_pretrained(COLBERT).to(device).eval()
check = colbert_tokenizer(questions[:4], padding=True, return_tensors="pt").to(device)
padding = check["attention_mask"] == 0
with torch.no_grad():
    theirs = reference(**check).last_hidden_state
    ours = colbert.encoder(check["input_ids"], padding)
gap = (ours - theirs)[~padding].abs().max().item()
print(f"ColBERTv2's encoder, largest difference from Hugging Face: {gap:.1e}")
assert gap < 1e-3
del reference


@torch.no_grad()
def colbert_queries(texts, length=32):
    batch = colbert_tokenizer([". " + t for t in texts], padding="max_length",
                              truncation=True, max_length=length,
                              return_tensors="pt").to(device)
    ids, attended = batch["input_ids"], batch["attention_mask"]
    ids[:, 1] = Q_MARKER  # the original ColBERT code's trick: ". " held the marker's place
    ids[ids == colbert_tokenizer.pad_token_id] = colbert_tokenizer.mask_token_id
    with torch.autocast(device, dtype=torch.float16, enabled=device == "cuda"):
        return colbert(ids, attended == 0).to(half), ids


@torch.no_grad()
def colbert_documents(texts, length=180):
    batch = colbert_tokenizer([". " + t for t in texts], padding=True, truncation=True,
                              max_length=length, return_tensors="pt").to(device)
    ids, attended = batch["input_ids"], batch["attention_mask"]
    ids[:, 1] = D_MARKER
    with torch.autocast(device, dtype=torch.float16, enabled=device == "cuda"):
        vectors = colbert(ids, attended == 0)
    keep = (attended == 1) & ~torch.isin(ids, PUNCTUATION)
    return (vectors * keep.unsqueeze(-1)).to(half), keep, ids

ColBERTv2's encoder, largest difference from Hugging Face: 2.1e-06

Before indexing anything, here is MaxSim at work on the first question and the paper it wants. Each row is a token of the question, each column a token of the abstract; the brightest cell of a row is the match that row contributes.

q_vectors, q_ids = colbert_queries(questions[:1])
d_vectors, d_keep, d_ids = colbert_documents([texts[gold[0][0]]])
q_tokens = colbert_tokenizer.convert_ids_to_tokens(q_ids[0])
d_tokens = colbert_tokenizer.convert_ids_to_tokens(d_ids[0])
shown = [j for j in range(len(d_tokens)) if d_keep[0, j]][:36]
n_words = q_tokens.index("[SEP]") + 1
similarity = (q_vectors[0, :n_words].float() @ d_vectors[0, shown].float().T).cpu().numpy()
heatmap(similarity, rows=q_tokens[:n_words], cols=[d_tokens[j] for j in shown], cell=0.26,
        vmin=0.0, vmax=1.0, title="MaxSim: the question against the paper's first tokens")
plt.show()
print("best match of each question token:")
print(", ".join(f"{q}→{d_tokens[shown[j]]}" for q, j in zip(q_tokens[2:n_words - 1],
                                                          similarity.argmax(1)[2:-1])))

MaxSim: the question against the paper

best match of each question token:
are→the, there→are, any→the, research→natural, papers→[unused1], on→of, methods→of, to→of, com→pre, ##press→##sti, large→tiny, -→of, scale→models, language→language, models→models, using→of, task→tasks, -→[CLS], ag→natural, ##nostic→bert, knowledge→understanding, di→di, ##sti→##sti, ##llation→##lling, techniques→of, ?→models

Most question tokens find their namesakes, piece by piece: di, ##sti, ##llation find the di, ##sti, ##lling of TinyBERT: Distilling BERT, and language and models find themselves.

Some matches are about meaning rather than spelling: large finds tiny. The question’s filler words, are, there, on, to, settle for whatever filler is nearest; they add a little to every paper’s score and decide nothing.

Now the whole corpus. The index keeps up to 180 vectors of 128 numbers for every paper. On a GPU we can afford to search it exhaustively: for every question, MaxSim against every paper, in blocks.

start = time.time()
order = np.argsort([-len(t) for t in texts])
index = torch.zeros(len(texts), 180, 128, dtype=half, device=device)
n_vectors = 0
for s in range(0, len(texts), 128):
    rows = order[s:s + 128]
    vectors, keep, _ = colbert_documents([texts[i] for i in rows])
    index[torch.as_tensor(rows, device=device), :vectors.shape[1]] = vectors
    n_vectors += keep.sum().item()
print(f"indexed in {time.time() - start:.0f} s: {n_vectors:,} vectors, "
      f"{n_vectors / len(texts):.0f} per paper")


@torch.no_grad()
def maxsim_search(query_vectors, block=2048):
    scores = torch.empty(len(query_vectors), len(index), device=device)
    for qs in range(0, len(query_vectors), 16):
        q = query_vectors[qs:qs + 16]
        for ds in range(0, len(index), block):
            sim = torch.einsum("aqd,nld->anql", q, index[ds:ds + block])
            scores[qs:qs + 16, ds:ds + block] = sim.amax(-1).float().sum(-1)
    return scores


start = time.time()
colbert_scores = maxsim_search(colbert_queries(questions)[0])
if device == "cuda":
    torch.cuda.synchronize()
evaluate("ColBERTv2", colbert_scores, time.time() - start, n_vectors * 128 * 2)

indexed in 209 s: 7,982,607 vectors, 138 per paper
methodrecall@5recall@20ms per questionindex, MB
ColBERTv20.470.6213.772,043.55

ColBERTv2 lands between the small and the large single-vector models: 0.47, above BM25 and all-MiniLM-L6-v2, below bge and the adapted MiniLM. Its price is the index, eight million vectors and 2 GB in float16, more than twenty times bge’s, and 14 ms per question even on a GPU.

Like the reranker, it learned what relevant means from MS MARCO’s web queries, not from questions about papers.

8. A search engine

Each method has seen the questions its own way, and their mistakes differ. BM25 cannot know that compressing a model and distillation are related; an embedding can miss that the question names a specific dataset.

The simplest way to combine rankings is reciprocal rank fusion[12]: each paper gets \(\sum_r 1 / (60 + \text{rank}_r)\) over the rankings \(r\) that contain it. It needs no training and no calibration of scores, which is the point: BM25’s scores and cosines are on scales that have nothing to do with each other, and ranks are on the same scale by construction.

A search engine in three stages. All papers go to two retrievers, BM25 on words and a bi-encoder on meaning, each returning its top 100; reciprocal rank fusion merges the two rankings; a cross-encoder reads the query with each of the 50 best candidates and returns the final top five A search engine in three stages. All papers go to two retrievers, BM25 on words and a bi-encoder on meaning, each returning its top 100; reciprocal rank fusion merges the two rankings; a cross-encoder reads the query with each of the 50 best candidates and returns the final top five

def fuse(*rankings, k=60, depth=100):
    fused = torch.zeros_like(rankings[0])
    weight = 1 / (k + torch.arange(1, depth + 1, device=fused.device, dtype=fused.dtype))
    for scores in rankings:
        top = scores.topk(depth, dim=1).indices
        fused.scatter_add_(1, top, weight.expand_as(top).contiguous())
    return fused


bm25_gpu = bm25_scores.to(device)
for name, scores in [("MiniLM-L6, trained here", student_scores),
                     ("all-MiniLM-L6-v2", minilm_scores),
                     ("all-MiniLM-L6-v2, adapted here", adapted_scores),
                     ("bge-base-en-v1.5", bge_scores)]:
    evaluate(f"BM25 + {name}", fuse(bm25_gpu, scores), 0, 0)
hybrid = fuse(bm25_gpu, adapted_scores)
reranked = rerank("BM25 + MiniLM adapted → reranker", hybrid)
engine_scores = fuse(hybrid, reranked, depth=50)  # the reranker's order and the hybrid's
evaluate("BM25 + MiniLM adapted + reranker, fused", engine_scores, 0, 0)
table(pd.DataFrame(results).drop(columns=["ms per question", "index, MB"]))
methodrecall@5recall@20ms per questionindex, MB
BM25 + MiniLM adapted → reranker0.550.71132.620.00
methodrecall@5recall@20
BM250.450.60
MiniLM-L6, as pre-trained0.000.00
MiniLM-L6, trained here0.320.46
all-MiniLM-L6-v20.420.59
all-MiniLM-L6-v2, [CLS] instead of the mean0.300.43
all-MiniLM-L6-v2, adapted here0.510.65
bge-base-en-v1.50.500.67
bge-base-en-v1.5, mean pooling, no instruction0.510.66
all-MiniLM-L6-v2 → reranker0.540.64
ColBERTv20.470.62
BM25 + MiniLM-L6, trained here0.470.63
BM25 + all-MiniLM-L6-v20.530.68
BM25 + all-MiniLM-L6-v2, adapted here0.560.70
BM25 + bge-base-en-v1.50.550.69
BM25 + MiniLM adapted → reranker0.550.71
BM25 + MiniLM adapted + reranker, fused0.590.73

Fusion is the best deal in the table. BM25 + all-MiniLM-L6-v2 reaches 0.53, eleven points above the model alone and eight above BM25, from two rankings that each find what the other misses. Even our ten-minute model, at 0.32 on its own, lifts BM25 by two points.

With the adapted MiniLM the hybrid reaches 0.56, as much as bge with BM25. The reranker on top of that hybrid changes nothing at five: on candidates this good, its judgement, learned from Bing queries, is no better than the hybrid’s.

Fused with the hybrid’s order, the two judgements together give the best result of the lesson, 0.59 at five and 0.73 at twenty.

Fusion and reranking add no index of their own, so their cost in the table is the cost of their parts. The whole picture, quality against the time a question takes and the memory its index needs:

frame = pd.DataFrame(results).set_index("method")
single = ["BM25", "MiniLM-L6, trained here", "all-MiniLM-L6-v2",
          "all-MiniLM-L6-v2, adapted here", "bge-base-en-v1.5", "ColBERTv2"]
points = {name: frame.loc[name] for name in single}
parts = frame.loc[["BM25", "all-MiniLM-L6-v2, adapted here",
                   "BM25 + MiniLM adapted → reranker"]].sum()
engine = "BM25 + MiniLM adapted + reranker, fused"
points["the search engine"] = pd.Series({"recall@5": frame.loc[engine, "recall@5"],
                                         "ms per question": parts["ms per question"],
                                         "index, MB": parts["index, MB"]})

fig, (left, right) = plt.subplots(1, 2, figsize=(9.6, 3.8), sharey=True)
for name, row in points.items():
    special = {"BM25": COLORS["data"], "the search engine": COLORS["att"]}
    color = special.get(name, COLORS["enc"])
    for ax, x in ((left, row["ms per question"]), (right, row["index, MB"])):
        ax.scatter(x, row["recall@5"], s=55, color=color, zorder=3)
        left_side = name in ("all-MiniLM-L6-v2, adapted here", "BM25")
        ax.annotate(name, (x, row["recall@5"]), xytext=(-6 if left_side else 6, 4),
                    textcoords="offset points", fontsize=8,
                    ha="right" if left_side else "left")
for ax in (left, right):
    ax.set_xscale("log")
    ax.margins(x=0.3)
left.set_xlabel("milliseconds per question, log scale")
right.set_xlabel("index size, MB, log scale")
left.set_ylabel("recall@5")
left.set_title("Quality against time")
right.set_title("Quality against memory")
plt.show()

Quality against time

The search engine is alone at the top of both charts. Almost all of its time is the reranker’s, about 130 ms per question on a T4; everything else answers in a millisecond or less, except ColBERT, which also needs twenty times the memory of anything else here.

Where the methods disagree is more instructive than the averages. For every question we find the rank of its first right paper under BM25 and under the adapted MiniLM, and look at the questions where one of them finds the paper at once and the other buries it.

def first_rank(scores):
    order = scores.argsort(dim=1, descending=True).cpu().numpy()
    position = np.empty_like(order)
    position[np.arange(len(order))[:, None], order] = np.arange(order.shape[1])
    return np.array([min(position[i, g]) + 1 for i, g in enumerate(gold)])


ranks = pd.DataFrame({"question": questions, "BM25": first_rank(bm25_gpu),
                      "MiniLM": first_rank(adapted_scores),
                      "engine": first_rank(engine_scores)})
words_win = ranks[(ranks["BM25"] <= 3) & (ranks["MiniLM"] > 50)]
meaning_wins = ranks[(ranks["MiniLM"] <= 3) & (ranks["BM25"] > 50)]
print(f"BM25 in the top 3, MiniLM below 50: {len(words_win)} questions; "
      f"the other way round: {len(meaning_wins)}")
display(table(pd.concat([words_win.head(3), meaning_wins.head(3)])))

BM25 in the top 3, MiniLM below 50: 13 questions; the other way round: 22
questionBM25MiniLMengine
In the context of natural language processing, I am looking for research that explores the relationship between a model's prediction entropy and its tendencies to copy existing text versus generating novel content. Can you recommend a paper?11623
How to achieve zero-shot lip reading?11874
Is there any work that allows large numbers of model outputs to be encoded and compared by causal language models in a single forward pass?38463
Can you recommend a conversational QA dataset where the human questioner does not have access to the evidence passage to simulate a more real-world information-seeking environment?155223
Can you suggest any literature that explores the idea of training neural networks to translate text passages into related questions?10014
When using pretrained transformer models for generating sentence embeddings, I've heard different strategies such as mean pooling and using the CLS token's embedding. What study shows that mean pooling outperforms CLS in semantic similarity tasks?59316

BM25 wins when a question hangs on a rare, exact term: zero-shot lip reading, prediction entropy, a single forward pass. The adapted model wins on paraphrase: the question about networks that “translate text passages into related questions” wants Learning to Ask: Neural Question Generation for Reading Comprehension, whose words the question doesn’t use.

Fusion keeps most of both, though not all: QuAC falls from second place under MiniLM to 23rd in the engine, dragged down by BM25’s 155th.

The last question in the table is about this lesson. It asks which study showed that mean pooling beats [CLS] for sentence similarity, and the paper it wants is Sentence-BERT: 59th for BM25, 3rd for the adapted MiniLM.

9. Your own search engine

Everything is in place for the practical result of the lesson: a function that takes a question in plain words and returns papers. Let’s ask it about the papers this course is built on.

@torch.no_grad()
def search(question, k=5, depth=50):
    words_ranked = torch.from_numpy(bm25.scores(question)).to(device)[None]
    meaning = embed(adapted, [question], 64) @ adapted_docs.T
    candidates = fuse(words_ranked, meaning).topk(depth, dim=1).indices[0].tolist()
    batch = reranker_tokenizer([question] * depth, [texts[j] for j in candidates],
                               padding=True, truncation="only_second", max_length=320,
                               return_tensors="pt").to(device)
    padding = batch["attention_mask"] == 0
    logits = reranker(batch["input_ids"], padding, batch["token_type_ids"])
    # fuse the two orders, as above: the candidates' own and the reranker's
    by_reranker = logits.argsort(descending=True).argsort()
    by_hybrid = torch.arange(depth, device=device)
    score = 1 / (61 + by_hybrid) + 1 / (61 + by_reranker)
    best = score.topk(k).indices.tolist()
    return [re.sub(r"\s+", " ", titles[candidates[j]])[:110] for j in best]


for question in [
    "a pre-training task where a discriminator spots tokens replaced by a small generator",
    "attention that keeps separate vectors for a token's content and its relative position",
    "do contextual word vectors of BERT occupy a narrow cone of the space?",
    "fast sentence similarity search with siamese BERT networks",
    "late interaction retrieval with compressed token vectors",
]:
    print(question)
    for rank, title in enumerate(search(question), 1):
        print(f"   {rank}. {title}")

a pre-training task where a discriminator spots tokens replaced by a small generator
   1. ELECTRA: PRE-TRAINING TEXT ENCODERS AS DISCRIMINATORS RATHER THAN GENERATORS
   2. Learning to Sample Replacements for ELECTRA Pre-Training
   3. ELECTRA: PRE-TRAINING TEXT ENCODERS AS DISCRIMINATORS RATHER THAN GENERATORS
   4. Training ELECTRA Augmented with Multi-word Selection
   5. Published as a conference paper at ICLR 2022 PRETRAINING TEXT ENCODERS WITH ADVERSARIAL MIXTURE OF TRAINING SI
attention that keeps separate vectors for a token's content and its relative position
   1. Self-Attention with Relative Position Representations
   2. PermuteFormer: Efficient Relative Position Encoding for Long Sequences
   3. A Simple and Effective Positional Encoding for Transformers
   4. Published as a conference paper at ICLR 2017 FRUSTRATINGLY SHORT ATTENTION SPANS IN NEURAL LANGUAGE MODELING
   5. A Target Attention Model for Neural Machine Translation
do contextual word vectors of BERT occupy a narrow cone of the space?
   1. How does BERT capture semantics? A closer look at polysemous words
   2. Unsupervised Contextualized Document Representation
   3. Open Sesame: Getting Inside BERT's Linguistic Knowledge
   4. Sentence Selection Strategies for Distilling Word Embeddings from BERT
   5. OoMMix: Out-of-manifold Regularization in Contextual Embedding Space for Text Classification
fast sentence similarity search with siamese BERT networks
   1. Sentence-BERT: Sentence Embeddings using Siamese BERT-Networks
   2. Structure-aware Sentence Encoder in BERT-Based Siamese Network
   3. Efficient Nearest Neighbor Emotion Classification with BERT-whitening
   4. Semantic Textual Similarity with Siamese Neural Networks
   5. An Unsupervised Sentence Embedding Method by Mutual Information Maximization
late interaction retrieval with compressed token vectors
   1. ColBERTv2: Effective and Efficient Retrieval via Lightweight Late Interaction
   2. CITADEL: Conditional Token Interaction via Dynamic Lexical Routing for Efficient and Effective Multi-Vector Re
   3. LI-RAGE: Late Interaction Retrieval Augmented Generation with Explicit Signals for Open-Domain Table Question 
   4. SDR: Efficient Neural Re-ranking using Succinct Document Representation
   5. Inflate and Shrink: Enriching and Reducing Interactions for Fast Text-Image Retrieval

Three of five come back first. ELECTRA appears twice because the corpus holds the paper twice, a duplicate that the same embeddings would catch in a deduplication pass.

The DeBERTa question returns an earlier paper on relative positions, Shaw et al.’s, and other position encodings, but not DeBERTa itself. The question about the narrow cone misses Ethayarajh’s paper, which is in the corpus; the engine prefers papers on BERT’s semantics that share more of the question’s words.

10. Which one to use today

The lesson’s table is a small version of what retrieval benchmarks keep finding. BM25 is a baseline that refuses to lose by much, and it is never the right thing to throw away: it finds names, numbers, codes and rare terms that an embedding smears into their neighbourhood.

A bi-encoder catches paraphrase and meaning, and costs one vector per document. Fusing the two is nearly free and usually better than either.

A cross-encoder on the top few dozen buys more, paid for in latency on every query: here it lifted a weak retriever by twelve points, and a strong hybrid only once its order was fused with the hybrid’s.

ColBERT puts token-level matching into the first stage, and pays in memory; ColBERTv2’s residual compression cuts that index “by 6–10×”, and the PLAID engine searches it up to “7× on a GPU and 45× on a CPU” faster than vanilla ColBERTv2 [13].

Choosing an embedding model is now a matter of reading a leaderboard such as MTEB [14] and then checking on your own data, because the leaderboard’s tasks are not yours.

Since 2024 its top places have mostly gone not to encoders but to decoders, large language models such as NV-Embed and Qwen3-Embedding turned into embedders by the same contrastive recipe.

On LitSearch itself the best embedder in the paper was one of them, the 7-billion-parameter GritLM, with 74.8% recall@5 on the specific questions. The next lessons are about decoders.

Two results of the lesson are worth taking to your own projects. The largest gain from a single step came from ten minutes of training a general model on the corpus’s own titles and abstracts, with no labels and no queries; if your documents have that kind of structure, use it before you buy a bigger model.

And the model card is part of the model: the same weights read with the wrong pooling lost twelve points without a word of warning.

Sentence-BERT’s passport, and the family tree with this lesson’s branch:

Sentence-BERT's passport: published August 2019 by Reimers and Gurevych at TU Darmstadt; a bi-encoder, one encoder applied to each text separately; built on BERT or RoBERTa; one vector per text, the mean of its token vectors, compared by cosine; trained in 2019 on about a million sentence pairs from SNLI and MultiNLI, today contrastively with in-batch negatives on up to a billion pairs; relatives are the cross-encoder and ColBERT; good at semantic search, retrieval, clustering and duplicates; not for exact names and numbers Sentence-BERT's passport: published August 2019 by Reimers and Gurevych at TU Darmstadt; a bi-encoder, one encoder applied to each text separately; built on BERT or RoBERTa; one vector per text, the mean of its token vectors, compared by cosine; trained in 2019 on about a million sentence pairs from SNLI and MultiNLI, today contrastively with in-batch negatives on up to a billion pairs; relatives are the cross-encoder and ColBERT; good at semantic search, retrieval, clustering and duplicates; not for exact names and numbers

The family of the Transformer as this course covers it, with lesson 3, embeddings and search, highlighted and lessons 1 and 2 done The family of the Transformer as this course covers it, with lesson 3, embeddings and search, highlighted and lessons 1 and 2 done

References

  1. Reimers, N., & Gurevych, I. (2019). Sentence-BERT: sentence embeddings using Siamese BERT-networks. EMNLP-IJCNLP 2019. ↩
  2. Ajith, A., Xia, M., Chevalier, A., Goyal, T., Chen, D., & Gao, T. (2024). LitSearch: a retrieval benchmark for scientific literature search. EMNLP 2024. ↩
  3. Robertson, S. E., Walker, S., Jones, S., Hancock-Beaulieu, M., & Gatford, M. (1995). Okapi at TREC-3. The Third Text REtrieval Conference (TREC-3), NIST SP 500-226. ↩
  4. Wang, W., Wei, F., Dong, L., Bao, H., Yang, N., & Zhou, M. (2020). MiniLM: deep self-attention distillation for task-agnostic compression of pre-trained transformers. NeurIPS 2020. ↩
  5. Henderson, M., Al-Rfou, R., Strope, B., Sung, Y., Lukács, L., Guo, R., Kumar, S., Miklos, B., & Kurzweil, R. (2017). Efficient natural language response suggestion for Smart Reply. arXiv:1705.00652. ↩
  6. Karpukhin, V., Oğuz, B., Min, S., Lewis, P., Wu, L., Edunov, S., Chen, D., & Yih, W. (2020). Dense passage retrieval for open-domain question answering. EMNLP 2020. ↩
  7. van den Oord, A., Li, Y., & Vinyals, O. (2018). Representation learning with contrastive predictive coding. arXiv:1807.03748. ↩
  8. Xiao, S., Liu, Z., Zhang, P., Muennighoff, N., Lian, D., & Nie, J.-Y. (2023). C-Pack: packed resources for general Chinese embeddings. SIGIR 2024. ↩
  9. Nogueira, R., & Cho, K. (2019). Passage re-ranking with BERT. arXiv:1901.04085. ↩
  10. Khattab, O., & Zaharia, M. (2020). ColBERT: efficient and effective passage search via contextualized late interaction over BERT. SIGIR 2020. ↩
  11. Santhanam, K., Khattab, O., Saad-Falcon, J., Potts, C., & Zaharia, M. (2022). ColBERTv2: effective and efficient retrieval via lightweight late interaction. NAACL 2022. ↩
  12. Cormack, G. V., Clarke, C. L. A., & Büttcher, S. (2009). Reciprocal rank fusion outperforms Condorcet and individual rank learning methods. SIGIR 2009. ↩
  13. Santhanam, K., Khattab, O., Potts, C., & Zaharia, M. (2022). PLAID: an efficient engine for late interaction retrieval. CIKM 2022. ↩
  14. Muennighoff, N., Tazi, N., Magne, L., & Reimers, N. (2022). MTEB: massive text embedding benchmark. EACL 2023. ↩
  15. Ethayarajh, K. (2019). How contextual are contextualized word representations? Comparing the geometry of BERT, ELMo, and GPT-2 embeddings. EMNLP-IJCNLP 2019.
]]> Eva Koroleva From RoBERTa to ModernBERT: what seven years did to BERT 2026-03-20T11:00:00+00:00 2026-11-02T00:00:00+00:00 https://blog.korolevalab.com/blog/2026/transformer-architectures-02-modernbert Transformer architectures: from BERT to diffusion · lesson 2 of 15 · all lessons · run on Kaggle · open in Colab · code on GitHub

Nine months after BERT, a team at Facebook AI and the University of Washington trained it again. They kept the architecture, every layer and every head, and changed only how it was trained: longer, on more data, with bigger batches and without one of its two training games.

The abstract of their paper says what came out: BERT “was significantly undertrained, and can match or exceed the performance of every model published after it” [1].

That sentence sets the tone for the next five years of encoders. ELECTRA changed the training task, DeBERTa changed how attention sees positions, and only in December 2024 did ModernBERT rebuild the block itself, out of the parts that large language models had made standard in the meantime.

Four encoders after BERT, each with the one thing it changed: RoBERTa in 2019 the training recipe, ELECTRA in 2020 the training task, DeBERTa in 2020 the attention, ModernBERT in 2024 the whole block Four encoders after BERT, each with the one thing it changed: RoBERTa in 2019 the training recipe, ELECTRA in 2020 the training task, DeBERTa in 2020 the attention, ModernBERT in 2024 the whole block

In this lesson you will write each change yourself, check it against the real weights, and then put the four generations to work on a job engineers know well: sorting GitHub issues into bugs, feature requests, questions and documentation, for projects the models have never seen.

What you’ll do

  1. Load RoBERTa into the BERT of lesson 1 with two small changes, and see what byte-level BPE does with code.
  2. Write ELECTRA’s replaced-token detection and race it against masked language modelling on a T4.
  3. Write DeBERTa’s disentangled attention and learn why its check needs float64.
  4. Write ModernBERT: rotary positions, GeGLU, no biases, and layers that alternate local and global attention.
  5. Fine-tune BERT, RoBERTa, DeBERTa-v3 and ModernBERT to triage issues of unseen repositories, and measure what reading the whole issue is worth.

What you need: a GPU. On Kaggle, Settings → Accelerator → GPU; on Colab, Runtime → Change runtime type → T4 GPU. The notebook downloads four pre-trained encoders and the issues of thirteen open-source projects; on a T4 it runs in about an hour.

Setup

One small helper package, tfs, draws the course’s plots and token chips and knows where the notebook runs (Kaggle, Colab or your own machine). Everything the lesson teaches is written out in the notebook itself.

import importlib.util
import subprocess
import sys

TFS = "git+https://github.com/xmarva/transformers-from-scratch@main"
if importlib.util.find_spec("tfs") is None:
    subprocess.check_call([sys.executable, "-m", "pip", "install", "-q", TFS])

from tfs import env
from tfs.viz import table, token_chips, use_course_style

env.quiet()  # no progress bars or environment warnings in the outputs
SEED = 42
env.seed_everything(SEED)
COLORS = use_course_style()
WORK_DIR = env.work_dir()

import ast
import json
import math
import re
import time
import warnings

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import torch
import torch.nn as nn
import torch.nn.functional as F
from huggingface_hub import hf_hub_download
from IPython.display import display
from safetensors.torch import load_file
from tfs.viz import heatmap
from transformers import AutoModelForSequenceClassification, AutoTokenizer
# the reference implementations, for the checks
from transformers import DebertaV2Model, ModernBertModel, RobertaModel

# mixed precision skips a step when a gradient overflows; PyTorch reports it as a warning
warnings.filterwarnings("ignore", message="Detected call of `lr_scheduler.step")
SMOKE = env.smoke()  # CI runs the lesson on a small slice; you run it in full
device = "cuda" if torch.cuda.is_available() else "cpu"
assert device == "cuda" or SMOKE, "This lesson fine-tunes several encoders: use a GPU."
if device == "cuda":
    print("GPU:", torch.cuda.get_device_name(0))

GPU: Tesla T4

1. RoBERTa: the same BERT, trained properly

RoBERTa’s authors went through BERT’s training choices one at a time, and most of the gain came from the plainest ones. BERT trained with batches of 256 sequences; RoBERTa with 8,000. BERT read 16 GB of text; RoBERTa “over 160GB”, for longer.

The cleverer-sounding changes mattered less. BERT masked each sentence once, when the data was prepared; RoBERTa draws new masks every time a sequence is used, as the mask_tokens of lesson 1 does, and found this “comparable or slightly better”.

Dropping BERT’s second game, telling a real next sentence from a random one, “matches or slightly improves” the results, so the inputs became whole stretches of text instead of pairs.

The architecture is BERT’s. So our BERT from lesson 1 should take RoBERTa’s weights, and it does, with two changes that are left-overs of the code RoBERTa was trained with. Here is lesson 1’s model again, in one cell:

class Embeddings(nn.Module):
    def __init__(self, config):
        super().__init__()
        d = config["hidden_size"]
        self.words = nn.Embedding(config["vocab_size"], d)
        self.positions = nn.Embedding(config["max_position_embeddings"], d)
        self.segments = nn.Embedding(config["type_vocab_size"], d)
        self.norm = nn.LayerNorm(d, eps=config["layer_norm_eps"])
        self.dropout = nn.Dropout(config["hidden_dropout_prob"])
        # RoBERTa numbers positions from pad_token_id + 1: fairseq kept 0 and 1 for padding
        self.first = config["pad_token_id"] + 1 if config["model_type"] == "roberta" else 0

    def forward(self, ids):
        positions = torch.arange(ids.shape[1], device=ids.device) + self.first
        segments = torch.zeros_like(ids)  # one sentence per input
        x = self.words(ids) + self.positions(positions) + self.segments(segments)
        return self.dropout(self.norm(x))


class Layer(nn.Module):
    def __init__(self, config):
        super().__init__()
        d, eps = config["hidden_size"], config["layer_norm_eps"]
        self.heads = config["num_attention_heads"]
        self.q, self.k, self.v, self.out = (nn.Linear(d, d) for _ in range(4))
        self.norm1 = nn.LayerNorm(d, eps=eps)
        self.up = nn.Linear(d, config["intermediate_size"])
        self.down = nn.Linear(config["intermediate_size"], d)
        self.norm2 = nn.LayerNorm(d, eps=eps)
        self.dropout = nn.Dropout(config["hidden_dropout_prob"])

    def split(self, x):
        batch, length, d = x.shape
        return x.view(batch, length, self.heads, d // self.heads).transpose(1, 2)

    def forward(self, x, padding):
        q, k, v = self.split(self.q(x)), self.split(self.k(x)), self.split(self.v(x))
        scores = q @ k.transpose(-2, -1) / math.sqrt(q.shape[-1])
        lowest = torch.finfo(scores.dtype).min  # not -inf: see section 4
        scores = scores.masked_fill(padding[:, None, None, :], lowest)
        y = (self.dropout(scores.softmax(dim=-1)) @ v).transpose(1, 2).reshape(x.shape)
        x = self.norm1(x + self.dropout(self.out(y)))
        return self.norm2(x + self.dropout(self.down(F.gelu(self.up(x)))))


class Encoder(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.embeddings = Embeddings(config)
        n_layers = config["num_hidden_layers"]
        self.layers = nn.ModuleList(Layer(config) for _ in range(n_layers))

    def forward(self, ids, padding):
        x = self.embeddings(ids)
        for layer in self.layers:
            x = layer(x, padding)
        return x

The second left-over is in the embedding table’s shape: RoBERTa has one segment, not two, since it never packs two sentences for the abandoned game, and 514 positions, two of them never used. The names of its tensors differ from Google’s too; they follow PyTorch’s conventions, weight and bias, because RoBERTa was trained in PyTorch.

def roberta_names(config):
    names = {
        "embeddings.words.weight": "roberta.embeddings.word_embeddings.weight",
        "embeddings.positions.weight": "roberta.embeddings.position_embeddings.weight",
        "embeddings.segments.weight": "roberta.embeddings.token_type_embeddings.weight",
        "embeddings.norm.weight": "roberta.embeddings.LayerNorm.weight",
        "embeddings.norm.bias": "roberta.embeddings.LayerNorm.bias",
    }
    parts = {"q": "attention.self.query", "k": "attention.self.key",
             "v": "attention.self.value", "out": "attention.output.dense",
             "norm1": "attention.output.LayerNorm", "up": "intermediate.dense",
             "down": "output.dense", "norm2": "output.LayerNorm"}
    for i in range(config["num_hidden_layers"]):
        for ours, theirs in parts.items():
            for p in ("weight", "bias"):
                names[f"layers.{i}.{ours}.{p}"] = f"roberta.encoder.layer.{i}.{theirs}.{p}"
    return names


ROBERTA = "FacebookAI/roberta-base"
with open(hf_hub_download(ROBERTA, "config.json")) as f:
    roberta_config = json.load(f)
weights = load_file(hf_hub_download(ROBERTA, "model.safetensors"))
roberta = Encoder(roberta_config)
state = {ours: weights[theirs] for ours, theirs in roberta_names(roberta_config).items()}
roberta.load_state_dict(state, strict=True)
roberta = roberta.to(device).eval()

roberta_tokenizer = AutoTokenizer.from_pretrained(ROBERTA)
reference = RobertaModel.from_pretrained(ROBERTA, add_pooling_layer=False).to(device).eval()
texts = ["Segfault in torch.compile() with dynamic shapes 🚀", "Docs typo"]
check = roberta_tokenizer(texts, padding=True, return_tensors="pt").to(device)
padding = check["attention_mask"] == 0
with torch.no_grad():
    theirs = reference(**check).last_hidden_state
    ours = roberta(check["input_ids"], padding)
gap = (ours - theirs)[~padding].abs().max().item()
print(f"largest difference from Hugging Face's RoBERTa: {gap:.1e}")
assert gap < 1e-3
del reference

largest difference from Hugging Face's RoBERTa: 5.7e-06

RoBERTa’s other visible change is its tokenizer. BERT’s WordPiece splits text into characters it knows, and a character it doesn’t know becomes [UNK], the same token for an emoji, a Chinese name and a Greek letter.

RoBERTa borrowed GPT-2’s byte-level BPE: the text is first turned into bytes, and any byte sequence can be spelled out with 256 byte tokens, so nothing is ever unknown.

The Ġ you’ll see marks a token that starts with a space, and pairs like ï are the two bytes of one character, ï, printed one by one.

bert_tokenizer = AutoTokenizer.from_pretrained("google-bert/bert-base-uncased")
samples = ["x = torch.compile(model)  # 🚀", "naïve café in Zürich, Ελλάδα"]
rows, labels = [], []
for text in samples:
    for name, tok in [("WordPiece", bert_tokenizer), ("byte BPE", roberta_tokenizer)]:
        rows.append(tok.tokenize(text))
        labels.append(name)
token_chips(rows, labels=labels, max_width=640,
            strong=lambda row, col, token: token == "[UNK]")
for name, tok in [("WordPiece", bert_tokenizer), ("byte BPE", roberta_tokenizer)]:
    unknown = sum(t == tok.unk_token for text in samples for t in tok.tokenize(text))
    print(f"{name}: {unknown} unknown tokens")

WordPiece: 1 unknown tokens
byte BPE: 0 unknown tokens

Output of the code above

WordPiece also lowercased everything and stripped the accents; byte BPE keeps the text exactly, and turns what it hasn’t seen often into more, smaller pieces, sometimes single bytes of a multi-byte character. For text that mixes prose, code, stack traces and emoji, which is what GitHub issues are, that matters.

2. ELECTRA: learning from every token

Masked language modelling has a cost that lesson 1 mentioned in passing: the model learns only from the 15% of tokens it hides. ELECTRA, from Stanford and Google, turned the game around [2]. A small masked language model, the generator, fills the masked positions with its own guesses.

The model we actually want, the discriminator, sees the result and says for every token whether it is the original or a replacement. This is replaced token detection, and its headline result was a small model trained “on one GPU for 4 days that outperforms GPT (trained using 30x more compute)”.

Replaced token detection. A sentence has some tokens masked; a small generator fills each mask with a plausible word, sometimes the right one; the discriminator reads the result and labels every token as original or replaced, so it learns from all of them Replaced token detection. A sentence has some tokens masked; a small generator fills each mask with a plausible word, sometimes the right one; the discriminator reads the result and labels every token as original or replaced, so it learns from all of them

The generator is trained as an ordinary masked language model, “rather than being trained adversarially to fool the discriminator”, and it is small: the paper recommends a quarter to half of the discriminator’s size.

Its guesses are good enough to be hard to spot, and when it happens to guess the right word, that token counts as original. After pre-training the generator is thrown away.

To see what this buys we’ll pre-train two small encoders from scratch, one with each task, for the same number of minutes on the same GPU, and then compare them on the job of section 5.

The text they learn from is GitHub issues themselves, which we load now. The dataset holds the issues of seventeen large projects with their labels [3]; every project labels its issues its own way, so we map the labels to four kinds.

ISSUES = "open-index/open-github-issues"
TRAIN_REPOS = ["python/cpython", "rust-lang/rust", "microsoft/TypeScript",
               "kubernetes/kubernetes", "pingcap/tidb", "swiftlang/swift",
               "ClickHouse/ClickHouse", "golang/go"]
TEST_REPOS = ["moby/moby", "etcd-io/etcd", "facebook/react", "redis/redis", "vuejs/core"]
KINDS = ["bug", "feature", "question", "docs"]
RULES = [("docs", r"\bdoc(s|umentation)?\b"), ("question", r"\bquestion\b|\bsupport\b"),
         ("feature", r"feature|enhancement|suggestion|proposal|improvement"),
         ("bug", r"\bbug\b|crash|unexpected behaviour")]
SKIP = r"not a bug|working as intended|duplicate|spam|question-answered"


def kind_of(labels):
    found = set()
    for label in labels:
        low = label.lower()
        if re.search(SKIP, low):
            return None
        for name, pattern in RULES:
            if re.search(pattern, low):
                found.add(name)
    return found.pop() if len(found) == 1 else None  # one kind, or the issue is left out


def load_issues(repo):
    path = hf_hub_download(ISSUES, f"data/issues/{repo}/0.parquet", repo_type="dataset")
    frame = pd.read_parquet(path, columns=["is_pull_request", "title", "body", "labels"])
    frame = frame[frame["is_pull_request"].astype(str) == "False"]
    if SMOKE:
        frame = frame.head(3000)
    labels = []
    for x in frame["labels"]:  # a list, the string of a list, or "null" for no labels
        if isinstance(x, str):
            labels.append(ast.literal_eval(x) if x.startswith("[") else [])
        else:
            labels.append(list(x) if x is not None else [])
    frame = frame.assign(kind=[kind_of(ls) for ls in labels], repo=repo)
    frame = frame[frame["kind"].notna()]
    text = frame["title"].fillna("") + "\n\n" + frame["body"].fillna("")
    return frame.assign(text=text)[["repo", "kind", "text"]]


issues = pd.concat([load_issues(r) for r in TRAIN_REPOS + TEST_REPOS], ignore_index=True)
table(pd.crosstab(issues["repo"], issues["kind"], margins=True).reset_index())
repobugdocsfeaturequestionAll
ClickHouse/ClickHouse8,4294564,5703,36216,817
etcd-io/etcd8501836904492,172
facebook/react1,41705795012,497
golang/go03,1377,8994611,082
kubernetes/kubernetes14,6361,0346,0252,31624,011
microsoft/TypeScript11,1193166,7023,31521,452
moby/moby3,0545062,9521,1607,672
pingcap/tidb11,258116,95697219,197
python/cpython27,4736,94413,204047,621
redis/redis1891120090490
rust-lang/rust19,9659186,934027,817
swiftlang/swift14,161301,4692315,683
vuejs/core957941101,377
All113,50813,55558,59112,234197,888

The projects in TEST_REPOS are kept for the end: neither pre-training nor fine-tuning will see a single issue of theirs. Both pre-training runs read the issues of the training projects, cut into pieces of 128 tokens.

train_text = issues.loc[issues["repo"].isin(TRAIN_REPOS), "text"]
n_sample = min(len(train_text), 400 if SMOKE else 40000)
sample = train_text.sample(n=n_sample, random_state=SEED)
encoded = bert_tokenizer(list(sample), add_special_tokens=False, truncation=True,
                         max_length=2048)["input_ids"]
cls, sep = bert_tokenizer.cls_token_id, bert_tokenizer.sep_token_id
pieces = []
for ids in encoded:
    for start in range(0, max(1, len(ids) - 32), 126):
        piece = [cls] + ids[start:start + 126] + [sep]
        pieces.append(piece + [0] * (128 - len(piece)))
corpus = torch.tensor(pieces)
real_tokens = corpus.ne(0).sum().item()
print(f"{len(corpus):,} pieces of 128 tokens, {real_tokens / 1e6:.1f}M real tokens")

221,157 pieces of 128 tokens, 27.0M real tokens

Both encoders are the same small BERT, 4 layers and 256 wide, initialised as BERT was. The masked language model gets lesson 1’s head. ELECTRA’s generator is a one-layer encoder that shares the discriminator’s word embeddings, a quarter of its depth; its head predicts words, and the discriminator’s head predicts one number per token, replaced or not.

TINY = dict(vocab_size=len(bert_tokenizer), hidden_size=256, num_hidden_layers=4,
            num_attention_heads=4, intermediate_size=1024, max_position_embeddings=512,
            type_vocab_size=2, layer_norm_eps=1e-12, hidden_dropout_prob=0.1,
            attention_probs_dropout_prob=0.1, model_type="bert", pad_token_id=0)


def init_like_bert(module):
    if isinstance(module, (nn.Linear, nn.Embedding)):
        nn.init.normal_(module.weight, std=0.02)
    if isinstance(module, nn.Linear) and module.bias is not None:
        nn.init.zeros_(module.bias)
    if isinstance(module, nn.LayerNorm):
        nn.init.ones_(module.weight)
        nn.init.zeros_(module.bias)


class WordHead(nn.Module):  # lesson 1's masked LM head: a score for every word
    def __init__(self, d, words):
        super().__init__()
        self.dense, self.norm, self.words = nn.Linear(d, d), nn.LayerNorm(d), words
        self.bias = nn.Parameter(torch.zeros(words.num_embeddings))

    def forward(self, h):
        return self.norm(F.gelu(self.dense(h))) @ self.words.weight.T + self.bias


class ReplacedHead(nn.Module):  # ELECTRA's discriminator: one score per token
    def __init__(self, d):
        super().__init__()
        self.dense, self.score = nn.Linear(d, d), nn.Linear(d, 1)

    def forward(self, h):
        return self.score(F.gelu(self.dense(h))).squeeze(-1)


def mask_tokens(ids, generator, rate=0.15):
    special = (ids == cls) | (ids == sep) | (ids == 0)
    chosen = (torch.rand(ids.shape, generator=generator) < rate) & ~special
    roll = torch.rand(ids.shape, generator=generator)
    masked = ids.clone()
    masked[chosen & (roll < 0.8)] = bert_tokenizer.mask_token_id
    swap = chosen & (roll >= 0.8) & (roll < 0.9)
    masked[swap] = torch.randint(len(bert_tokenizer), ids.shape, generator=generator)[swap]
    return masked, chosen

One training step of each. The masked language model computes its loss at the chosen positions only. ELECTRA runs the generator on the masked text, samples a word for every chosen position from the generator’s own predictions, and asks the discriminator to find the replacements among all real tokens.

The paper weighs the discriminator’s loss fifty times more than the generator’s, since a binary decision per token gives much smaller losses than a choice among 30,000 words.

def mlm_step(models, ids, generator):
    encoder, head = models
    masked, chosen = mask_tokens(ids, generator)
    logits = head(encoder(masked.to(device), (ids == 0).to(device)))
    return F.cross_entropy(logits[chosen.to(device)].float(), ids[chosen].to(device))


def electra_step(models, ids, generator):
    encoder, head, gen, gen_head = models
    masked, chosen = mask_tokens(ids, generator)
    padding = (ids == 0).to(device)
    ids, masked, chosen = ids.to(device), masked.to(device), chosen.to(device)
    gen_logits = gen_head(gen(masked, padding))[chosen].float()
    gen_loss = F.cross_entropy(gen_logits, ids[chosen])
    with torch.no_grad():
        guesses = torch.distributions.Categorical(logits=gen_logits).sample()
    corrupted = ids.clone()
    corrupted[chosen] = guesses
    replaced = (corrupted != ids).float()  # a lucky guess counts as original
    scores = head(encoder(corrupted, padding)).float()
    disc_loss = F.binary_cross_entropy_with_logits(scores[~padding], replaced[~padding])
    return gen_loss + 50 * disc_loss


def pretrain(task, minutes):
    torch.manual_seed(SEED)
    encoder = Encoder(TINY)
    if task == "mlm":
        models = [encoder, WordHead(256, encoder.embeddings.words)]
    else:
        generator_config = dict(TINY, num_hidden_layers=1)
        gen = Encoder(generator_config)
        gen.embeddings.words = encoder.embeddings.words  # shared word embeddings
        models = [encoder, ReplacedHead(256), gen, WordHead(256, encoder.embeddings.words)]
    for model in models:
        model.apply(init_like_bert)
    models = [m.to(device).train() for m in models]
    params = {id(p): p for m in models for p in m.parameters()}.values()
    optimizer = torch.optim.AdamW(params, lr=5e-4, weight_decay=0.01)
    step_fn = mlm_step if task == "mlm" else electra_step
    generator = torch.Generator().manual_seed(SEED)
    scaler = torch.amp.GradScaler(enabled=device == "cuda")  # float16 gradients in range
    steps, start = 0, time.time()
    while time.time() - start < minutes * 60:
        for group in optimizer.param_groups:
            group["lr"] = 5e-4 * min(1.0, (steps + 1) / 500)  # warm-up, then constant
        batch = corpus[torch.randint(len(corpus), (128,), generator=generator)]
        with torch.autocast(device, dtype=torch.float16, enabled=device == "cuda"):
            loss = step_fn(models, batch, generator)
        optimizer.zero_grad(set_to_none=True)
        scaler.scale(loss).backward()
        scaler.unscale_(optimizer)
        nn.utils.clip_grad_norm_(params, 1.0)
        scaler.step(optimizer)
        scaler.update()
        steps += 1
    return encoder.eval(), steps


MINUTES = 0.3 if SMOKE else 8
mlm_encoder, mlm_steps = pretrain("mlm", MINUTES)
electra_encoder, electra_steps = pretrain("electra", MINUTES)
print(f"{MINUTES} minutes each: masked LM made {mlm_steps:,} steps, "
      f"ELECTRA {electra_steps:,}")

8 minutes each: masked LM made 2,661 steps, ELECTRA 2,143

ELECTRA’s step is more expensive, since it runs two encoders, so in the same minutes it makes fewer steps. What each pre-trained encoder learned shows in two ways.

First, without any training of the encoder: its vectors for whole issues, frozen, and a linear classifier on top that tells bugs from feature requests, questions and documentation.

Second, the way encoders are used in practice: fine-tuned, with only a hundred labelled issues of each kind, since pre-training matters most when labels are scarce. A third encoder starts from random weights. All the labelled issues come from the training projects, and the test issues from the five unseen ones.

def balanced(frame, per_kind, seed):
    parts = []
    for kind in KINDS:
        rows = frame[frame["kind"] == kind]
        parts.append(rows.sample(n=min(per_kind, len(rows)), random_state=seed))
    return pd.concat(parts).sample(frac=1, random_state=seed).reset_index(drop=True)


seen = issues[issues["repo"].isin(TRAIN_REPOS)]
unseen = issues[issues["repo"].isin(TEST_REPOS)]
probe_train = balanced(seen, 16 if SMOKE else 1000, SEED)
few_train = balanced(seen, 8 if SMOKE else 100, SEED + 2)
small_test = balanced(unseen, 8 if SMOKE else 300, SEED)
print(f"{len(probe_train)} issues for the probe, {len(few_train)} for fine-tuning, "
      f"{len(small_test)} test issues from unseen projects")


def macro_f1(truth, predicted):  # F1 of each kind, averaged: every kind counts the same
    truth, predicted = np.asarray(truth), np.asarray(predicted)
    scores = []
    for c in range(len(KINDS)):
        hits = np.sum((predicted == c) & (truth == c))
        precision = hits / max(1, np.sum(predicted == c))
        recall = hits / max(1, np.sum(truth == c))
        scores.append(0.0 if hits == 0 else 2 * precision * recall / (precision + recall))
    return float(np.mean(scores))


def confusion(truth, predicted):  # rows: the actual kind, columns: what the model said
    counts = np.zeros((len(KINDS), len(KINDS)))
    for t, p in zip(truth, predicted, strict=True):
        counts[t, p] += 1
    return counts / np.maximum(1, counts.sum(axis=1, keepdims=True))


def batches(frame, size, shuffle):
    order = np.random.permutation(len(frame)) if shuffle else np.arange(len(frame))
    for k in range(0, len(frame), size):
        rows = frame.iloc[order[k:k + size]]
        enc = bert_tokenizer(list(rows["text"]), truncation=True, max_length=128,
                             padding=True, return_tensors="pt")
        labels = torch.tensor([KINDS.index(x) for x in rows["kind"]])
        yield enc["input_ids"].to(device), labels.to(device)


def pooled(encoder, ids):  # the mean of the last layer over real tokens
    real = (ids != 0).unsqueeze(-1).float()
    return (encoder(ids, ids == 0) * real).sum(1) / real.sum(1)


@torch.no_grad()
def frozen_features(encoder, frame):
    encoder.eval()
    xs, ys = [], []
    for ids, labels in batches(frame, 64, shuffle=False):
        xs.append(pooled(encoder, ids))
        ys.append(labels)
    return torch.cat(xs), torch.cat(ys)


def probe_f1(encoder):
    x, y = frozen_features(encoder, probe_train)
    x_test, y_test = frozen_features(encoder, small_test)
    mean, std = x.mean(0), x.std(0) + 1e-6
    torch.manual_seed(SEED)
    linear = nn.Linear(x.shape[1], len(KINDS)).to(device)
    optimizer = torch.optim.Adam(linear.parameters(), lr=1e-2, weight_decay=1e-4)
    for _ in range(300):  # the whole set at once: it is small
        loss = F.cross_entropy(linear((x - mean) / std), y)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
    predicted = linear((x_test - mean) / std).argmax(-1)
    return macro_f1(y_test.tolist(), predicted.tolist())


def fine_tuned_f1(encoder, epochs=8):
    encoder = encoder.train()
    torch.manual_seed(SEED)
    head = nn.Linear(256, len(KINDS)).to(device)
    params = list(encoder.parameters()) + list(head.parameters())
    optimizer = torch.optim.AdamW(params, lr=1e-4, weight_decay=0.01)
    for _ in range(epochs):
        for ids, labels in batches(few_train, 32, shuffle=True):
            loss = F.cross_entropy(head(pooled(encoder, ids)), labels)
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
    encoder.eval()
    predicted, truth = [], []
    with torch.no_grad():
        for ids, labels in batches(small_test, 64, shuffle=False):
            predicted += head(pooled(encoder, ids)).argmax(-1).tolist()
            truth += labels.tolist()
    return macro_f1(truth, predicted)


untrained = Encoder(TINY)
untrained.apply(init_like_bert)
untrained = untrained.to(device)
race = []
for name, steps, encoder in [("none, random weights", 0, untrained),
                             ("masked LM", mlm_steps, mlm_encoder),
                             ("ELECTRA", electra_steps, electra_encoder)]:
    probe = probe_f1(encoder)  # before fine-tuning changes the weights
    fine = fine_tuned_f1(encoder)
    race.append({"pre-training": name, "steps": steps, "frozen + linear probe": probe,
                 "fine-tuned on 400 issues": fine})
table(pd.DataFrame(race))

4000 issues for the probe, 400 for fine-tuning, 1200 test issues from unseen projects
pre-trainingstepsfrozen + linear probefine-tuned on 400 issues
none, random weights00.470.30
masked LM2,6610.450.49
ELECTRA2,1430.470.32

The probe column first. Averaged vectors of an untrained encoder are a bag of random word vectors, and a linear classifier on top already does far better than chance: which words an issue contains says a lot about its kind. Eight minutes of pre-training didn’t make either encoder’s frozen vectors better than that.

Fine-tuning on few labels is where pre-training shows. The masked language model starts far ahead of random weights, the effect of lesson 1 at a tiny scale. ELECTRA, in our race, did not: its discriminator fine-tuned no better than random weights.

The paper’s smallest ELECTRA trained for four days on a GPU, some seven hundred times our budget, with a generator narrower than its discriminator, while ours is one layer of the same width; its advantage is measured there.

Key point

A race this short says what happens in the first eight minutes, and in them the simpler task won. It is RoBERTa’s lesson again: an idea is only as good as the recipe it is trained with, and ours is a sketch.

3. DeBERTa: positions as vectors of their own

BERT adds a position vector to each word vector at the input, and from then on a token is one vector that mixes what the word is with where it stands.

DeBERTa, from Microsoft, keeps them apart [4]. Positions are relative, the distance from one token to another, and each layer’s attention score between tokens \(i\) and \(j\) is the sum of three terms:

\[A_{ij} = \underbrace{q_i^{c} \cdot k_j^{c}}_{\text{content to content}} + \underbrace{q_i^{c} \cdot k_{\delta(i,j)}^{r}}_{\text{content to position}} + \underbrace{k_j^{c} \cdot q_{\delta(j,i)}^{r}}_{\text{position to content}}, \qquad \text{scaled by } \frac{1}{\sqrt{3d}} .\]

Here \(q^c, k^c\) are the usual queries and keys of the tokens, and \(q^r, k^r\) are queries and keys computed, with the same projection matrices, from a table of relative-position embeddings; \(\delta(i, j)\) is the distance from \(i\) to \(j\), put into a bucket. The second term asks “what does this word look for at that distance”, the third “what does that word look for at this distance”. The scale has 3 instead of 1 because three terms are summed.

DeBERTa's attention score between two tokens as the sum of three terms: content to content, as in BERT; content to position, the query of the first token against the embedding of their relative distance; and position to content, the key of the second token against the distance seen from the other side DeBERTa's attention score between two tokens as the sum of three terms: content to content, as in BERT; content to position, the query of the first token against the embedding of their relative distance; and position to content, the key of the second token against the distance seen from the other side

It worked: at the end of 2020 a 1.5-billion-parameter DeBERTa became the first model to score above the human baseline on SuperGLUE, a benchmark built to be hard for models, 89.9 against 89.8. Microsoft’s own announcement added that the model “is by no means reaching the human-level intelligence of NLU” [5].

DeBERTa-v3 sorts distances into 256 buckets: exact up to 128 tokens in either direction, and logarithmically coarser beyond, so a model trained on 512-token inputs can still tell near from far in longer ones:

def log_buckets(relative, buckets=256, max_position=512):
    mid = buckets // 2
    sign = torch.sign(relative)
    near = (relative < mid) & (relative > -mid)
    size = torch.where(near, torch.tensor(mid - 1), relative.abs())
    scale = math.log((max_position - 1) / mid)
    far = torch.ceil(torch.log(size / mid) / scale * (mid - 1)) + mid
    return torch.where(size <= mid, relative, (far * sign).long())


distance = torch.arange(-700, 701)
fig, ax = plt.subplots(figsize=(6.4, 3.4))
ax.plot(distance, log_buckets(distance), color=COLORS["att"])
ax.set_title("DeBERTa-v3's bucket for each relative distance")
ax.set_xlabel("distance from the query to the key, tokens")
ax.set_ylabel("bucket")
plt.show()

DeBERTa-v3

The rest of a DeBERTa-v3 layer is BERT’s: post-norm, GELU. One more detail of v3: there are no absolute positions anywhere, not even at the input.

class DebertaLayer(Layer):
    def forward(self, x, padding, rel, relative, span=256):
        q, k, v = self.split(self.q(x)), self.split(self.k(x)), self.split(self.v(x))
        pq, pk = self.split(self.q(rel)), self.split(self.k(rel))  # same projections
        n, heads = x.shape[1], self.heads
        shape = (x.shape[0], heads, n, n)
        c2p_index = (relative + span).clamp(0, 2 * span - 1).expand(shape)
        p2c_index = (-relative + span).clamp(0, 2 * span - 1).expand(shape)
        content = q @ k.transpose(-2, -1)
        c2p = torch.gather(q @ pk.transpose(-2, -1), -1, c2p_index)
        p2c = torch.gather(k @ pq.transpose(-2, -1), -1, p2c_index).transpose(-1, -2)
        scores = (content + c2p + p2c) / math.sqrt(3 * q.shape[-1])
        lowest = torch.finfo(scores.dtype).min
        scores = scores.masked_fill(padding[:, None, None, :], lowest)
        y = (self.dropout(scores.softmax(dim=-1)) @ v).transpose(1, 2).reshape(x.shape)
        x = self.norm1(x + self.dropout(self.out(y)))
        return self.norm2(x + self.dropout(self.down(F.gelu(self.up(x)))))


class Deberta(nn.Module):
    def __init__(self, config):
        super().__init__()
        d, eps = config["hidden_size"], config["layer_norm_eps"]
        self.span = config["position_buckets"]
        self.words = nn.Embedding(config["vocab_size"], d)
        self.norm = nn.LayerNorm(d, eps=eps)
        self.distances = nn.Embedding(2 * self.span, d)
        self.distance_norm = nn.LayerNorm(d, eps=eps)
        n_layers = config["num_hidden_layers"]
        self.layers = nn.ModuleList(DebertaLayer(config) for _ in range(n_layers))

    def forward(self, ids, padding):
        x = self.norm(self.words(ids)) * (~padding)[..., None]
        positions = torch.arange(ids.shape[1], device=ids.device)
        relative = log_buckets(positions[:, None] - positions[None, :], self.span)
        rel = self.distance_norm(self.distances.weight)[None]
        for layer in self.layers:
            x = layer(x, padding, rel, relative, self.span)
        return x

Microsoft published DeBERTa-v3’s weights as a PyTorch pickle rather than safetensors; weights_only=True loads the tensors and refuses anything else a pickle could carry.

DEBERTA = "microsoft/deberta-v3-base"
with open(hf_hub_download(DEBERTA, "config.json")) as f:
    deberta_config = json.load(f)
deberta_file = hf_hub_download(DEBERTA, "pytorch_model.bin")
deberta_weights = torch.load(deberta_file, map_location="cpu", weights_only=True)
names = {"words.weight": "deberta.embeddings.word_embeddings.weight",
         "norm.weight": "deberta.embeddings.LayerNorm.weight",
         "norm.bias": "deberta.embeddings.LayerNorm.bias",
         "distances.weight": "deberta.encoder.rel_embeddings.weight",
         "distance_norm.weight": "deberta.encoder.LayerNorm.weight",
         "distance_norm.bias": "deberta.encoder.LayerNorm.bias"}
parts = {"q": "attention.self.query_proj", "k": "attention.self.key_proj",
         "v": "attention.self.value_proj", "out": "attention.output.dense",
         "norm1": "attention.output.LayerNorm", "up": "intermediate.dense",
         "down": "output.dense", "norm2": "output.LayerNorm"}
for i in range(deberta_config["num_hidden_layers"]):
    for ours, theirs in parts.items():
        for p in ("weight", "bias"):
            names[f"layers.{i}.{ours}.{p}"] = f"deberta.encoder.layer.{i}.{theirs}.{p}"
deberta = Deberta(deberta_config)
deberta.load_state_dict({o: deberta_weights[t] for o, t in names.items()}, strict=True)
deberta.eval()

deberta_tokenizer = AutoTokenizer.from_pretrained(DEBERTA)
reference = DebertaV2Model.from_pretrained(DEBERTA).eval()
trace = "Traceback (most recent call last): File 'train.py', line 12 " * 25
check = deberta_tokenizer([trace], return_tensors="pt")
for dtype in (torch.float32, torch.float64):
    with torch.no_grad():
        theirs = reference.to(dtype)(**check).last_hidden_state
        ours = deberta.to(dtype)(check["input_ids"], check["attention_mask"] == 0)
    print(f"{str(dtype)[6:]}: largest difference {(ours - theirs).abs().max().item():.1e}, "
          f"largest activation {theirs.abs().max().item():.1f}")
del reference, deberta, deberta_weights

float32: largest difference 1.7e-05, largest activation 9.5
float64: largest difference 3.1e-07, largest activation 9.5

In float32 the two implementations agree to four or five decimal places, the rounding of summing the same numbers in a different order; in float64, to seven. Our forty lines are DeBERTa’s attention.

The explicit .to(dtype) in that loop is there for a reason we found the hard way while writing this lesson.

Trap

Microsoft stored DeBERTa-v3’s weights in float16, and recent versions of transformers load a model in the precision of its checkpoint. Compared without the cast, the reference quietly runs in half precision, and the two models disagree in the second decimal place, which looks exactly like a bug in the attention. Running both in float64 is the quickest way to tell a bug from rounding: a bug survives it, rounding doesn’t.

4. ModernBERT: an encoder rebuilt from LLM parts

For five years after DeBERTa, decoder-only models got all the attention and nearly all the improvements. ModernBERT, from Answer.AI and LightOn, took those improvements back into an encoder [6].

Every piece of its block has appeared in this course before, mostly in lesson 5 of the first one: pre-normalisation, rotary positions (RoPE) instead of learned ones, a gated feed-forward layer (GeGLU, the GELU sibling of Llama’s SwiGLU), and no bias terms except in the last layer that scores the vocabulary.

One piece is new: most layers attend only to a window of 128 tokens around each token, and every third layer attends globally. It was trained on 2 trillion tokens of text and code, with 30% of tokens masked rather than 15%, and its authors call it “a major Pareto improvement over older encoders”: better and faster at once.

ModernBERT-base's 22 layers: every third layer, starting with the first, attends globally with rotary positions of base 160,000; the others attend to a window of 128 tokens with rotary base 10,000. Each layer is pre-norm, without biases, with a GeGLU feed-forward network ModernBERT-base's 22 layers: every third layer, starting with the first, attends globally with rotary positions of base 160,000; the others attend to a window of 128 tokens with rotary base 10,000. Each layer is pre-norm, without biases, with a GeGLU feed-forward network

MODERNBERT = "answerdotai/ModernBERT-base"
with open(hf_hub_download(MODERNBERT, "config.json")) as f:
    modern_config = json.load(f)


def rotate(x, theta):
    d = x.shape[-1]
    inv = 1.0 / theta ** (torch.arange(0, d, 2, device=x.device, dtype=torch.float32) / d)
    angles = torch.arange(x.shape[-2], device=x.device, dtype=torch.float32)[:, None] * inv
    cos, sin = torch.cat([angles.cos()] * 2, -1), torch.cat([angles.sin()] * 2, -1)
    half = d // 2
    return x * cos + torch.cat([-x[..., half:], x[..., :half]], -1) * sin


class ModernLayer(nn.Module):
    def __init__(self, config, index):
        super().__init__()
        d, eps = config["hidden_size"], config["norm_eps"]
        self.heads = config["num_attention_heads"]
        self.is_global = index % config["global_attn_every_n_layers"] == 0
        self.theta = config["global_rope_theta" if self.is_global else "local_rope_theta"]
        self.half_window = config["local_attention"] // 2
        # the first layer reads the embeddings' norm directly, so it has none of its own
        if index > 0:
            self.attn_norm = nn.LayerNorm(d, eps=eps, bias=False)
        else:
            self.attn_norm = nn.Identity()
        self.qkv = nn.Linear(d, 3 * d, bias=False)
        self.out = nn.Linear(d, d, bias=False)
        self.mlp_norm = nn.LayerNorm(d, eps=eps, bias=False)
        self.up = nn.Linear(d, 2 * config["intermediate_size"], bias=False)
        self.down = nn.Linear(config["intermediate_size"], d, bias=False)

    def attention(self, x, padding):
        b, n, d = x.shape
        qkv = self.qkv(x).view(b, n, 3, self.heads, d // self.heads).permute(2, 0, 3, 1, 4)
        q, k, v = rotate(qkv[0], self.theta), rotate(qkv[1], self.theta), qkv[2]
        hidden = padding[:, None, None, :]
        if not self.is_global:
            position = torch.arange(n, device=x.device)
            distance = (position[:, None] - position[None, :]).abs()
            hidden = hidden | (distance > self.half_window)
        scores = q @ k.transpose(-2, -1) / math.sqrt(q.shape[-1])
        weights = scores.masked_fill(hidden, torch.finfo(scores.dtype).min).softmax(-1)
        return self.out((weights @ v).transpose(1, 2).reshape(b, n, d))

    def forward(self, x, padding):
        x = x + self.attention(self.attn_norm(x), padding)
        h, gate = self.up(self.mlp_norm(x)).chunk(2, dim=-1)  # GeGLU
        return x + self.down(F.gelu(h) * gate)


class ModernBert(nn.Module):
    def __init__(self, config):
        super().__init__()
        d, eps = config["hidden_size"], config["norm_eps"]
        self.words = nn.Embedding(config["vocab_size"], d)
        self.norm = nn.LayerNorm(d, eps=eps, bias=False)
        n_layers = config["num_hidden_layers"]
        self.layers = nn.ModuleList(ModernLayer(config, i) for i in range(n_layers))
        self.final_norm = nn.LayerNorm(d, eps=eps, bias=False)

    def forward(self, ids, padding):
        x = self.norm(self.words(ids))
        for layer in self.layers:
            x = layer(x, padding)
        return self.final_norm(x)

A padded position far from any real token has every key hidden in a local layer, both by padding and by distance. A softmax over nothing but \(-\infty\) is nan, and nan times a zero weight is still nan, which would leak into the real tokens through the next layer. Filling hidden scores with the most negative finite number instead of \(-\infty\), as above, keeps such rows harmless.

weights = load_file(hf_hub_download(MODERNBERT, "model.safetensors"))
names = {"words.weight": "model.embeddings.tok_embeddings.weight",
         "norm.weight": "model.embeddings.norm.weight",
         "final_norm.weight": "model.final_norm.weight"}
for i in range(modern_config["num_hidden_layers"]):
    for ours, theirs in [("qkv", "attn.Wqkv"), ("out", "attn.Wo"), ("up", "mlp.Wi"),
                         ("down", "mlp.Wo"), ("mlp_norm", "mlp_norm")]:
        names[f"layers.{i}.{ours}.weight"] = f"model.layers.{i}.{theirs}.weight"
    if i > 0:
        names[f"layers.{i}.attn_norm.weight"] = f"model.layers.{i}.attn_norm.weight"
modern = ModernBert(modern_config)
modern.load_state_dict({o: weights[t] for o, t in names.items()}, strict=True)
modern = modern.to(device).eval()

modern_tokenizer = AutoTokenizer.from_pretrained(MODERNBERT)
reference = ModernBertModel.from_pretrained(MODERNBERT, attn_implementation="eager")
reference = reference.to(device).eval()
check = modern_tokenizer([trace, "Docs typo"], padding=True, return_tensors="pt").to(device)
padding = check["attention_mask"] == 0
with torch.no_grad():
    theirs = reference(**check).last_hidden_state
    ours = modern(check["input_ids"], padding)
gap = (ours - theirs)[~padding].abs().max().item()
length = check["input_ids"].shape[1]
print(f"{length} tokens, largest difference from Hugging Face's ModernBERT: {gap:.1e}")
assert gap < 1e-3
del reference

403 tokens, largest difference from Hugging Face's ModernBERT: 1.6e-04

What does the window buy? In a global layer each of \(n\) tokens attends to all \(n\); in a local layer to at most 129. Counted over ModernBERT-base’s 22 layers:

rows = []
for n in [512, 2048, 8192]:
    pairs = sum(n * n if layer.is_global else n * min(n, 2 * layer.half_window + 1)
                for layer in modern.layers)
    rows.append({"tokens": n, "attended pairs": f"{pairs / 1e6:,.0f}M",
                 "share of an all-global model": f"{pairs / (22 * n * n):.0%}"})
table(pd.DataFrame(rows))
tokensattended pairsshare of an all-global model
5123M52%
2,04837M40%
8,192552M37%

The local layers get almost free as inputs grow, and the cost settles at that of the eight global layers: 37% of an all-global model at 8,192 tokens, the length ModernBERT was trained for.

On GPUs with FlashAttention the saving is real time; the T4 has no kernel for windows, so here the mask only hides scores that are still computed.

5. Triage for projects the model has never seen

Every large project gets more issues than its maintainers can read the day they arrive, and the first thing anyone does with a new issue is decide what it is: a bug, a feature request, a question, or a problem with the documentation.

A model that does this well saves real time, and the honest test of one is a project it was not trained on, since that is where a new tool would be used.

The four generations get the same job: 6,000 labelled issues from the eight training projects, 1,500 of each kind, and the issues of the five unseen projects as the test, 400 of each kind. Each model is the pre-trained encoder with a classification head, fine-tuned for two epochs with the same recipe, reading the first 256 tokens of each issue.

train_set = balanced(seen, 8 if SMOKE else 1500, SEED + 1)
test_set = balanced(unseen, 4 if SMOKE else 400, SEED + 1)
TOKENS, EPOCHS = (64, 1) if SMOKE else (256, 2)  # CI checks the code on a CPU, briefly
lengths = [len(modern_tokenizer(t)["input_ids"]) for t in test_set["text"]]
print(f"test issues longer than 256 tokens: {np.mean(np.array(lengths) > 256):.0%}, "
      f"longer than 1,024: {np.mean(np.array(lengths) > 1024):.0%}")
example = test_set[test_set["kind"] == "bug"].iloc[0]
print(example["repo"], "|", example["kind"], "\n", example["text"][:700])

test issues longer than 256 tokens: 56%, longer than 1,024: 21%
moby/moby | bug 
 docker push fails when DNS `AAAA` response with 0 answers beats `A` response with 1 answer

I've been experiencing an issue with docker push for a local, ipv4 only registry. Every once in a while, it'll fail, with `dial tcp: lookup registry.weave.local: no such host`, only to succeed after.

I've run wireshark and have confirmed that Docker is issuing an `A` and `AAAA` DNS query simultaneously, and when the `AAAA` response comes back first, I get the error. Most of the time, the `A` response comes back first, and it succeeds.

Also, I've confirmed that AAAA is responding correctly with "No answer".

![Screen shot from wireshark showing the proper AAAA response](https://files.gitter.im/docker
ENCODERS = {"BERT": "google-bert/bert-base-uncased", "RoBERTa": ROBERTA,
            "DeBERTa-v3": DEBERTA, "ModernBERT": MODERNBERT}


def fine_tune(name, max_length, label, epochs=EPOCHS, lr=3e-5, batch_size=16):
    torch.manual_seed(SEED)
    tokenizer = AutoTokenizer.from_pretrained(ENCODERS[name])
    model = AutoModelForSequenceClassification.from_pretrained(ENCODERS[name],
                                                               num_labels=len(KINDS))
    model = model.float().to(device)  # some checkpoints are stored in float16
    steps = epochs * math.ceil(len(train_set) / batch_size)
    optimizer = torch.optim.AdamW(model.parameters(), lr=lr, weight_decay=0.01)
    warmup = max(1, steps // 10)

    def rate(step):
        return min((step + 1) / warmup, max(0.0, (steps - step) / max(1, steps - warmup)))

    schedule = torch.optim.lr_scheduler.LambdaLR(optimizer, rate)
    scaler = torch.amp.GradScaler(enabled=device == "cuda")

    def batches(frame, size, shuffle):
        order = np.random.permutation(len(frame)) if shuffle else np.arange(len(frame))
        for k in range(0, len(frame), size):
            rows = frame.iloc[order[k:k + size]]
            enc = tokenizer(list(rows["text"]), truncation=True, max_length=max_length,
                            padding=True, return_tensors="pt").to(device)
            labels = torch.tensor([KINDS.index(x) for x in rows["kind"]], device=device)
            yield enc, labels

    t0 = time.time()
    model.train()
    for _ in range(epochs):
        for enc, labels in batches(train_set, batch_size, shuffle=True):
            with torch.autocast(device, dtype=torch.float16, enabled=device == "cuda"):
                loss = F.cross_entropy(model(**enc).logits.float(), labels)
            optimizer.zero_grad(set_to_none=True)
            scaler.scale(loss).backward()
            scaler.unscale_(optimizer)
            nn.utils.clip_grad_norm_(model.parameters(), 1.0)
            scaler.step(optimizer)
            scaler.update()
            schedule.step()
    train_minutes = (time.time() - t0) / 60
    model.eval()
    predicted, truth = [], []
    t0 = time.time()
    half = torch.autocast(device, dtype=torch.float16, enabled=device == "cuda")
    with torch.no_grad(), half:
        for enc, labels in batches(test_set, 32, shuffle=False):
            predicted += model(**enc).logits.argmax(-1).tolist()
            truth += labels.tolist()
    per_second = len(test_set) / (time.time() - t0)
    result = {"model": label, "tokens": max_length,
              "macro F1": macro_f1(truth, predicted),
              "train, min": train_minutes, "issues/s": per_second}
    del model
    if device == "cuda":
        torch.cuda.empty_cache()
    return result, confusion(truth, predicted)


results, confusions = [], {}
for name in ENCODERS:
    result, confusions[name] = fine_tune(name, TOKENS, label=name)
    results.append(result)
    print(f"{name}: macro F1 {result['macro F1']:.3f}", flush=True)
table(pd.DataFrame(results))

BERT: macro F1 0.703
RoBERTa: macro F1 0.717
DeBERTa-v3: macro F1 0.728
ModernBERT: macro F1 0.730
modeltokensmacro F1train, minissues/s
BERT2560.702.85223.35
RoBERTa2560.722.92220.92
DeBERTa-v32560.735.16119.90
ModernBERT2560.734.93113.90

The table needs a second look before any conclusion: each number comes from one run, one random seed, and a test set of 1,600 issues, so differences of a point or two are within the noise.

With that in mind, the order is the order of history. RoBERTa gains on BERT with exactly the same block, from its training alone; DeBERTa-v3 and ModernBERT add a little more and tie with each other.

On this GPU they also cost the same, about half the speed of BERT; ModernBERT’s speed comes from FlashAttention and unpadding, which a T4 can’t use.

Now the thing only ModernBERT can do here. BERT, RoBERTa and DeBERTa-v3 were trained on 512 tokens at most, while ModernBERT was trained on 8,192. Issues are long, with logs and stack traces, so we let it read four times more:

LONG = f"ModernBERT, {4 * TOKENS:,} tokens"
long_result, confusions[LONG] = fine_tune("ModernBERT", 4 * TOKENS, LONG, batch_size=8)
results.append(long_result)
display(table(pd.DataFrame(results)))

fig, ax = plt.subplots(figsize=(6.4, 3.8))
for row in results:
    color = COLORS["att"] if row["model"].startswith("ModernBERT") else COLORS["enc"]
    ax.scatter(row["issues/s"], row["macro F1"], s=60, color=color)
    ax.annotate(row["model"], (row["issues/s"], row["macro F1"]), xytext=(6, 4),
                textcoords="offset points", fontsize=9)
ax.set_title("Triage of unseen projects: quality against speed")
ax.set_xlabel("issues classified per second on this GPU")
ax.set_ylabel("macro F1")
plt.show()
modeltokensmacro F1train, minissues/s
BERT2560.702.85223.35
RoBERTa2560.722.92220.92
DeBERTa-v32560.735.16119.90
ModernBERT2560.734.93113.90
ModernBERT, 1,024 tokens1,0240.7319.6426.31

Triage of unseen projects: quality against speed

best = max(results, key=lambda row: row["macro F1"])["model"]
heatmap(confusions[best], rows=KINDS, cols=KINDS, annotate=True, fmt="{:.2f}", cell=0.6,
        title=f"{best}: what each kind of issue was taken for", xlabel="predicted",
        ylabel="actual", role="enc")
plt.show()

Output of the code above

Reading four times more of each issue changed nothing: ModernBERT trained four times longer and classified four times slower for the same macro F1, 0.73.

The title and the first paragraph decide what an issue is, and the stack trace below them adds nothing to that decision. Long context earns its price where the answer is spread through a document, like the personal data of lesson 1’s bank documents, not here.

The errors that remain gather in one place: almost a third of the bugs were taken for questions. A bug report often is a question, “is this expected?”, and the projects’ own labels don’t agree on where one ends and the other begins.

6. Which one to use today

Encoders are not a museum piece. On the day this lesson was written, the four most downloaded models on the Hugging Face Hub were small BERT-style encoders that turn text into vectors for search, the first of them more than a hundred million times in a month[7]. The next lesson is about exactly those.

Key point

For a new project that needs an encoder, ModernBERT is a good default: it reads long inputs, it was trained on code as well as prose, and it is fast where FlashAttention is available.

DeBERTa-v3 remains a strong choice for short English text when the last point of accuracy matters more than speed, and years of Kaggle competitions have worn smooth every way of fine-tuning it.

RoBERTa and BERT are worth knowing because half the models on the Hub are their descendants, fine-tuned for one task or another, and an existing fine-tuned model often beats a better base trained by you on less data.

The lesson’s other result is about method rather than models. RoBERTa won without touching the architecture, ELECTRA without touching it either, and our tiny race showed that the training task alone moves the result. When a paper reports a new architecture beating an old one, the first question is whether the old one was trained as well.

ModernBERT’s passport, and the family tree with this lesson’s branch:

ModernBERT's passport: published December 2024 by Answer.AI and LightOn; encoder; masked language model with 30% of tokens masked, on 2 trillion tokens of text and code; attention alternating global and local windows of 128 tokens; rotary positions, 8,192 tokens; pre-norm LayerNorm, no biases except the vocabulary decoder; GeGLU; base 149 million parameters and 22 layers, large 395 million and 28 layers ModernBERT's passport: published December 2024 by Answer.AI and LightOn; encoder; masked language model with 30% of tokens masked, on 2 trillion tokens of text and code; attention alternating global and local windows of 128 tokens; rotary positions, 8,192 tokens; pre-norm LayerNorm, no biases except the vocabulary decoder; GeGLU; base 149 million parameters and 22 layers, large 395 million and 28 layers

The family of the Transformer as this course covers it, with lesson 2, RoBERTa to ModernBERT, highlighted and lesson 1, BERT, done The family of the Transformer as this course covers it, with lesson 2, RoBERTa to ModernBERT, highlighted and lesson 1, BERT, done

References

  1. Liu, Y., Ott, M., Goyal, N., Du, J., Joshi, M., Chen, D., Levy, O., Lewis, M., Zettlemoyer, L., & Stoyanov, V. (2019). RoBERTa: a robustly optimized BERT pretraining approach. arXiv:1907.11692. ↩
  2. Clark, K., Luong, M.-T., Le, Q. V., & Manning, C. D. (2020). ELECTRA: pre-training text encoders as discriminators rather than generators. ICLR 2020. ↩
  3. open-index (2026). Open GitHub issues. Hugging Face dataset, ODC-By. ↩
  4. He, P., Liu, X., Gao, J., & Chen, W. (2020). DeBERTa: decoding-enhanced BERT with disentangled attention. ICLR 2021. ↩
  5. Microsoft Research, 2021. microsoft.com. ↩
  6. Warner, B., Chaffin, A., Clavié, B., Weller, O., Hallström, O., Taghadouini, S., Gallagher, A., Biswas, R., Ladhak, F., Aarsen, T., Cooper, N., Adams, G., Howard, J., & Poli, I. (2024). Smarter, better, faster, longer: a modern bidirectional encoder for fast, memory efficient, and long context finetuning and inference. arXiv:2412.13663. ↩
  7. Hugging Face, 2026. huggingface.co. ↩
  8. Devlin, J., Chang, M.-W., Lee, K., & Toutanova, K. (2018). BERT: pre-training of deep bidirectional transformers for language understanding. NAACL 2019.
  9. He, P., Gao, J., & Chen, W. (2021). DeBERTaV3: improving DeBERTa using ELECTRA-style pre-training with gradient-disentangled embedding sharing. ICLR 2023.
  10. Su, J., Lu, Y., Pan, S., Murtadha, A., Wen, B., & Liu, Y. (2021). RoFormer: enhanced transformer with rotary position embedding. arXiv:2104.09864.
  11. Shazeer, N. (2020). GLU variants improve Transformer. arXiv:2002.05202.
]]> Eva Koroleva BERT from scratch: masked language modelling and fine-tuning 2026-02-27T10:00:00+00:00 2026-11-02T00:00:00+00:00 https://blog.korolevalab.com/blog/2026/transformer-architectures-01-bert Transformer architectures: from BERT to diffusion · lesson 1 of 15 · all lessons · run on Kaggle · open in Colab · code on GitHub

In October 2019 Google changed the way its search engine reads queries. Its own example was the query “2019 brazil traveler to usa need a visa”. The old system treated the little word to as noise and “returned results about U.S. citizens traveling to Brazil”, which is the other way around.

The new one understood that the traveller goes from Brazil to the US. Google said the model would help with “one in 10 searches in the U.S. in English” and called it “the biggest leap forward in the past five years” [1].

The query 2019 brazil traveler to usa need a visa, read two ways. As a bag of keywords, brazil, traveler, usa and visa match pages for Americans going to Brazil. Read as a sentence, the word to links traveler and usa: a Brazilian going to the US The query 2019 brazil traveler to usa need a visa, read two ways. As a bag of keywords, brazil, traveler, usa and visa match pages for Americans going to Brazil. Read as a sentence, the word to links traveler and usa: a Brazilian going to the US

The model was BERT, published by Google a year earlier [2]. It is the Transformer encoder you built in the first course, trained on books and Wikipedia to fill in hidden words, and then taught a task with a few thousand examples.

That recipe, pre-train once and fine-tune for everything, is how most of NLP worked for the next four years, and BERT-sized encoders still sort, tag and search a large part of the world’s text, because they are small, fast and very good at it.

In this lesson you will write BERT yourself, load Google’s weights into your code and check that it gives exactly the same answers as the reference implementation. Then you will ask it what it learned, look at what its attention heads do, and teach it a job that companies pay for: finding personal data in financial documents.

What you’ll do

  1. See why a language model couldn’t simply read in both directions, and how a mask fixed it.
  2. Read BERT-base’s blueprint from its config and count where its 110 million parameters live.
  3. Write BERT in PyTorch, load Google’s weights into it and match Hugging Face’s outputs.
  4. Mask words the way BERT was trained, and ask it about facts, grammar, negation and gender.
  5. Look inside its attention: which heads watch the previous word, and why so many watch [SEP].
  6. Fine-tune it to find names, addresses and account numbers in long documents, and compare it with regular expressions and with the same network trained without pre-training.

What you need: a GPU. On Kaggle, Settings → Accelerator → GPU; on Colab, Runtime → Change runtime type → T4 GPU.

The notebook downloads BERT-base twice (uncased and cased, 440 MB each) and a dataset of synthetic financial documents, and fine-tunes seven models, most of them on a few hundred documents; on a T4 it runs in about an hour.

Setup

One small helper package, tfs, draws the course’s plots and token chips and knows where the notebook runs (Kaggle, Colab or your own machine). Everything the lesson teaches is written out in the notebook itself.

import importlib.util
import subprocess
import sys

TFS = "git+https://github.com/xmarva/transformers-from-scratch@main"
if importlib.util.find_spec("tfs") is None:
    subprocess.check_call([sys.executable, "-m", "pip", "install", "-q", TFS])

from tfs import env
from tfs.viz import table, token_chips, use_course_style

env.quiet()  # no progress bars or environment warnings in the outputs
SEED = 42
env.seed_everything(SEED)
COLORS = use_course_style()
WORK_DIR = env.work_dir()

import json
import math
import random
import re
import time

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import torch
import torch.nn as nn
import torch.nn.functional as F
from datasets import load_dataset
from huggingface_hub import get_safetensors_metadata, hf_hub_download
from IPython.display import HTML, display
from safetensors.torch import load_file
from tfs.viz import heatmap
from transformers import AutoTokenizer, BertForMaskedLM

import warnings

# mixed precision skips a step when a gradient overflows; PyTorch reports it as a warning
warnings.filterwarnings("ignore", message="Detected call of `lr_scheduler.step")
SMOKE = env.smoke()  # CI runs the lesson on a small slice; you run it in full
device = "cuda" if torch.cuda.is_available() else "cpu"
assert device == "cuda" or SMOKE, "This lesson fine-tunes BERT seven times: use a GPU."
if device == "cuda":
    print("GPU:", torch.cuda.get_device_name(0))

GPU: Tesla T4

1. Reading in both directions

Before BERT, the standard way to give a network some knowledge of language was a table of word vectors, such as word2vec’s: one vector per word, learned from which words appear near which. One vector per word is the problem.

Bank gets the same vector in “the bank raised its rates” and “we sat on the river bank”, and the network that uses it has to work out the difference on its own, from whatever little labelled data it has.

The fix everyone could see was to read the whole sentence first: a vector for bank computed from its neighbours. A language model does exactly that.

It is trained on raw text, which is free and endless, to predict the next word, so its hidden state at every position is a summary of the words before it. In 2018 two models showed how much that summary is worth.

ELMo trained two recurrent language models, one reading left to right and one right to left, and glued their states together[3]. OpenAI’s first GPT trained a Transformer decoder left to right and fine-tuned the whole of it for each task [4].

Both read with one eye closed. In “we sat on the river bank” the word that decides what bank means comes before it, but in “the bank of the river was steep” it comes after, and a left-to-right model sees it too late.

ELMo’s two directions never meet inside the network: each half reads one way, and only their outputs are joined.

Why not train one model that reads both ways at once? Because it would cheat. Devlin and colleagues put it in one sentence:

Bidirectional conditioning would allow each word to indirectly “see itself”, and the model could trivially predict the target word in a multi-layered context.

Devlin et al., 2018

Ask a model to predict word 5 while it can attend to word 5, and it learns to copy. Even if word 5 can’t attend to itself in the first layer, word 6 can, and in the second layer word 5 reads it back from word 6.

Three ways to learn from raw text. GPT predicts each word from the words before it. ELMo trains a left-to-right and a right-to-left model separately and joins their outputs. BERT hides a word behind a mask and predicts it from both sides at once, in every layer Three ways to learn from raw text. GPT predicts each word from the words before it. ELMo trains a left-to-right and a right-to-left model separately and joins their outputs. BERT hides a word behind a mask and predicts it from both sides at once, in every layer

BERT’s answer is old. In 1953 Wilson Taylor proposed measuring how readable a text is by deleting every fifth word and asking people to fill the gaps; he called it the cloze procedure [5].

BERT plays the same game. It hides some of the words, replaces them with a special token [MASK] and is trained to restore them.

A hidden word can’t see itself, so every layer can attend to everything else, in both directions. The paper calls this the masked language model (MLM), “inspired by the Cloze task (Taylor, 1953)”.

The price is efficiency. A left-to-right model learns from every word of every sentence; BERT learns only from the 15% it hides.

Google paid for that with compute: BERT-base was trained “on 4 Cloud TPUs in Pod configuration (16 TPU chips total)” and “each pre-training took 4 days to complete”, on a corpus of 3.3 billion words, books and English Wikipedia. We will not repeat that. We will write the network and borrow the four days.

2. The blueprint

BERT-base’s weights live on the Hugging Face Hub, under Google’s name. Its configuration is a few hundred bytes:

BERT = "google-bert/bert-base-uncased"
with open(hf_hub_download(BERT, "config.json")) as f:
    config = json.load(f)

FIELDS = ["num_hidden_layers", "hidden_size", "num_attention_heads", "intermediate_size",
          "hidden_act", "max_position_embeddings", "type_vocab_size", "vocab_size",
          "layer_norm_eps"]
table(pd.DataFrame({"setting": FIELDS, "value": [str(config[name]) for name in FIELDS]}))
settingvalue
num_hidden_layers12
hidden_size768
num_attention_heads12
intermediate_size3072
hidden_actgelu
max_position_embeddings512
type_vocab_size2
vocab_size30522
layer_norm_eps1e-12

Twelve layers, each 768 wide with 12 attention heads, a feed-forward layer four times wider, GELU instead of ReLU.

This is the encoder of the 2017 Transformer with two differences you will see in the code: positions are learned vectors rather than sinusoids, and there is a third embedding, the segment, which says whether a token belongs to the first or the second of two sentences.

512 positions is the most BERT can read at once, a limit we will run into in section 7.

The headers of the weight file list every tensor with its shape, and the Hub serves just the headers. Summed by where they live:

metadata = get_safetensors_metadata(BERT)
tensors = {}
for file_metadata in metadata.files_metadata.values():
    for name, info in file_metadata.tensors.items():
        tensors[name] = info.shape


def part_of(name):
    if name == "bert.embeddings.word_embeddings.weight":
        return "word embeddings"
    if name.startswith("bert.embeddings."):
        return "positions, segments, norm"
    if ".attention." in name:
        return "attention, 12 layers"
    if ".intermediate." in name or ".output." in name:
        return "feed-forward, 12 layers"
    if name.startswith("bert.pooler."):
        return "pooler"
    return "pre-training heads"


parts = {}
for name, shape in tensors.items():
    parts[part_of(name)] = parts.get(part_of(name), 0) + math.prod(shape)
total = sum(parts.values())
print(f"{len(tensors)} tensors, {total:,} parameters")

fig, ax = plt.subplots(figsize=(6.4, 2.9))
names = list(parts)[::-1]
ax.barh(names, [parts[n] / 1e6 for n in names], color=COLORS["enc"])
for k, n in enumerate(names):
    label = f"{parts[n] / 1e6:.1f}M"
    ax.text(parts[n] / 1e6 + 0.8, k, label, va="center", color=COLORS["muted"])
ax.set_title("Where BERT-base's parameters live")
ax.set_xlabel("millions of parameters")
ax.grid(axis="y", visible=False)
plt.show()

206 tensors, 110,106,428 parameters

Where BERT-base

One matrix stands out. The word embeddings, one 768-number row for each of the 30,522 entries of the vocabulary, take a fifth of the model, as much as three of its twelve layers. That is why the multilingual BERT, which needs a vocabulary for 104 languages, is bigger than the English one with exactly the same layers.

Two names in the list are a little archaeological. Every LayerNorm stores gamma and beta rather than PyTorch’s weight and bias: Google trained BERT in TensorFlow, and the converted weights kept their names. And there is no output matrix for the masked word prediction at all.

The model predicts a word by comparing its hidden state with the very same embedding table it reads words from; the input and output vocabularies share one matrix.

3. WordPiece and what goes in

BERT reads WordPiece tokens. WordPiece builds its vocabulary like BPE in lesson 1 of the first course, by merging frequent pieces, but it picks the merge that most increases the likelihood of the training text rather than the most frequent pair, and it marks the pieces that continue a word with ##. Here are two sentences as BERT receives them:

tokenizer = AutoTokenizer.from_pretrained(BERT)
pair = tokenizer("The bank raised its rates.", "We sat on the river bank, unbothered.")
tokens = tokenizer.convert_ids_to_tokens(pair["input_ids"])
token_chips([tokens], ids=[pair["input_ids"]], mark_spaces=False, max_width=620,
            strong=lambda row, col, token: token in ("[CLS]", "[SEP]"))
print("segments:", pair["token_type_ids"])

segments: [0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]

Output of the code above

[CLS] opens every input; its final hidden state is meant to summarise the whole input for classification. [SEP] closes each sentence.

The segment ids, 0 for the first sentence and 1 for the second, were needed for BERT’s second training game, next sentence prediction: half of the training pairs were real consecutive sentences and half were random ones, and BERT had to tell which.

The game turned out to be nearly useless, and the next lesson’s RoBERTa dropped it. Unbothered was not frequent enough to get a token of its own and is spelled out in pieces: nothing is ever out of vocabulary.

This is the uncased BERT: its tokenizer lowercases everything and strips accents before splitting. Google also released a cased one. The difference matters more than it sounds:

cased_tokenizer = AutoTokenizer.from_pretrained("google-bert/bert-base-cased")
name = "Eufemia Gordillo-Pagès"
token_chips([tokenizer.tokenize(name), cased_tokenizer.tokenize(name)],
            labels=["uncased", "cased"], max_width=620)

Output of the code above

Output of the code above

Trap

Lowercasing, which helps a model that reads queries typed in a hurry, hurts a model that must find names: the capital letters are gone, and stripping the accent turns the end of a Spanish surname into the English word pages.

In section 7 we will look for names, and switch to the cased BERT.

Each token becomes the sum of three learned vectors, its word, its position and its segment:

BERT's input for two sentences: each token's vector is the sum of a word embedding, a learned position embedding for positions 0 to 511, and a segment embedding A or B, followed by LayerNorm BERT's input for two sentences: each token's vector is the sum of a word embedding, a learned position embedding for positions 0 to 511, and a segment embedding A or B, followed by LayerNorm

4. BERT in PyTorch

The embeddings first: three tables, added, normalised. Dropout is part of the model; it is active while training and does nothing in eval() mode.

class Embeddings(nn.Module):
    def __init__(self, config):
        super().__init__()
        d = config["hidden_size"]
        self.words = nn.Embedding(config["vocab_size"], d)
        self.positions = nn.Embedding(config["max_position_embeddings"], d)
        self.segments = nn.Embedding(config["type_vocab_size"], d)
        self.norm = nn.LayerNorm(d, eps=config["layer_norm_eps"])
        self.dropout = nn.Dropout(config["hidden_dropout_prob"])

    def forward(self, ids, segments):
        positions = torch.arange(ids.shape[1], device=ids.device)
        x = self.words(ids) + self.positions(positions) + self.segments(segments)
        return self.dropout(self.norm(x))

Self-attention is the one from the first course, with one difference that is the whole point: there is no causal mask. Every token attends to every other token. The only mask hides padding, the filler that makes the sentences of a batch equally long.

padding is True at those positions, and their scores become \(-\infty\) before the softmax, so they get zero weight. The layer can also hand back its attention weights, which we will look at in section 6.

class SelfAttention(nn.Module):
    def __init__(self, config):
        super().__init__()
        d = config["hidden_size"]
        self.heads = config["num_attention_heads"]
        self.q = nn.Linear(d, d)
        self.k = nn.Linear(d, d)
        self.v = nn.Linear(d, d)
        self.out = nn.Linear(d, d)
        self.dropout = nn.Dropout(config["attention_probs_dropout_prob"])

    def split(self, x):
        batch, length, d = x.shape
        return x.view(batch, length, self.heads, d // self.heads).transpose(1, 2)

    def forward(self, x, padding, keep_weights=False):
        q, k, v = self.split(self.q(x)), self.split(self.k(x)), self.split(self.v(x))
        scores = q @ k.transpose(-2, -1) / math.sqrt(q.shape[-1])
        scores = scores.masked_fill(padding[:, None, None, :], float("-inf"))
        weights = scores.softmax(dim=-1)
        y = self.dropout(weights) @ v
        y = y.transpose(1, 2).reshape(x.shape)
        return self.out(y), (weights if keep_weights else None)

A layer of BERT normalises after adding the residual, as the 2017 Transformer did; the next lessons’ models moved the norm before. The feed-forward part uses GELU, \(\text{GELU}(x) = x\,\Phi(x)\), where \(\Phi\) is the standard normal distribution function: a ReLU whose corner is smoothed, so that small negative inputs pass a little.

One BERT layer: self-attention, dropout, add the input, LayerNorm; then a feed-forward network 768 to 3072 to 768 with GELU, dropout, add, LayerNorm. The norm sits after each addition, on the main path One BERT layer: self-attention, dropout, add the input, LayerNorm; then a feed-forward network 768 to 3072 to 768 with GELU, dropout, add, LayerNorm. The norm sits after each addition, on the main path

class Layer(nn.Module):
    def __init__(self, config):
        super().__init__()
        d, eps = config["hidden_size"], config["layer_norm_eps"]
        self.attention = SelfAttention(config)
        self.norm1 = nn.LayerNorm(d, eps=eps)
        self.up = nn.Linear(d, config["intermediate_size"])
        self.down = nn.Linear(config["intermediate_size"], d)
        self.norm2 = nn.LayerNorm(d, eps=eps)
        self.dropout = nn.Dropout(config["hidden_dropout_prob"])

    def forward(self, x, padding, keep_weights=False):
        a, weights = self.attention(x, padding, keep_weights)
        x = self.norm1(x + self.dropout(a))
        f = self.down(F.gelu(self.up(x)))
        x = self.norm2(x + self.dropout(f))
        return x, weights


class Bert(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.embeddings = Embeddings(config)
        n_layers = config["num_hidden_layers"]
        self.layers = nn.ModuleList(Layer(config) for _ in range(n_layers))

    def forward(self, ids, segments, padding, keep_weights=False):
        x = self.embeddings(ids, segments)
        attention = []
        for layer in self.layers:
            x, weights = layer(x, padding, keep_weights)
            attention.append(weights)
        return x, attention

The head that predicts hidden words is small: a dense layer, GELU, LayerNorm, and then a dot product of the result with every row of the word embedding table, plus a bias per word. The table is not copied; the head holds the same nn.Embedding the model reads with, so training the head trains the embeddings too.

class MaskedLMHead(nn.Module):
    def __init__(self, config, words):
        super().__init__()
        d = config["hidden_size"]
        self.dense = nn.Linear(d, d)
        self.norm = nn.LayerNorm(d, eps=config["layer_norm_eps"])
        self.words = words  # the input embeddings, shared
        self.bias = nn.Parameter(torch.zeros(config["vocab_size"]))

    def forward(self, h):
        h = self.norm(F.gelu(self.dense(h)))
        return h @ self.words.weight.T + self.bias


class MaskedLM(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.bert = Bert(config)
        self.head = MaskedLMHead(config, self.bert.embeddings.words)

    def forward(self, ids, segments, padding):
        h, _ = self.bert(ids, segments, padding)
        return self.head(h)

Now the transplant. Our names differ from Google’s, so we write down which of their tensors goes where, load the file and let load_state_dict refuse anything missing or left over:

def google_names(config):
    names = {
        "bert.embeddings.words.weight": "bert.embeddings.word_embeddings.weight",
        "bert.embeddings.positions.weight": "bert.embeddings.position_embeddings.weight",
        "bert.embeddings.segments.weight": "bert.embeddings.token_type_embeddings.weight",
        "bert.embeddings.norm.weight": "bert.embeddings.LayerNorm.gamma",
        "bert.embeddings.norm.bias": "bert.embeddings.LayerNorm.beta",
        "head.dense.weight": "cls.predictions.transform.dense.weight",
        "head.dense.bias": "cls.predictions.transform.dense.bias",
        "head.norm.weight": "cls.predictions.transform.LayerNorm.gamma",
        "head.norm.bias": "cls.predictions.transform.LayerNorm.beta",
        "head.words.weight": "bert.embeddings.word_embeddings.weight",  # the shared table
        "head.bias": "cls.predictions.bias",
    }
    layer_parts = {
        "attention.q": "attention.self.query", "attention.k": "attention.self.key",
        "attention.v": "attention.self.value", "attention.out": "attention.output.dense",
        "up": "intermediate.dense", "down": "output.dense",
    }
    norms = {"norm1": "attention.output.LayerNorm", "norm2": "output.LayerNorm"}
    for i in range(config["num_hidden_layers"]):
        ours, theirs = f"bert.layers.{i}.", f"bert.encoder.layer.{i}."
        for a, b in layer_parts.items():
            names[ours + a + ".weight"] = theirs + b + ".weight"
            names[ours + a + ".bias"] = theirs + b + ".bias"
        for a, b in norms.items():
            names[ours + a + ".weight"] = theirs + b + ".gamma"
            names[ours + a + ".bias"] = theirs + b + ".beta"
    return names


def load_masked_lm(repo):
    with open(hf_hub_download(repo, "config.json")) as f:
        config = json.load(f)
    weights = load_file(hf_hub_download(repo, "model.safetensors"))
    model = MaskedLM(config)
    state = {ours: weights[theirs] for ours, theirs in google_names(config).items()}
    model.load_state_dict(state, strict=True)
    return model, config


mlm, config = load_masked_lm(BERT)
mlm = mlm.to(device).eval()
print(f"loaded {sum(p.numel() for p in mlm.parameters()):,} parameters into our model")

loaded 109,514,298 parameters into our model

Two of Google’s tensors were not used: the pooler and the next-sentence head, which only the abandoned second game needs. Does our model compute what Google’s does?

Hugging Face’s BertForMaskedLM is the reference implementation. Both models get the same two sentences; we compare the hidden states of the last layer and the scores of every word of the vocabulary at every position:

reference = BertForMaskedLM.from_pretrained(BERT).to(device).eval()
check = tokenizer(["The bank raised its rates.", "We sat on the river bank, unbothered."],
                  padding=True, return_tensors="pt").to(device)
padding = check["attention_mask"] == 0
with torch.no_grad():
    theirs = reference(**check, output_hidden_states=True)
    hidden, _ = mlm.bert(check["input_ids"], check["token_type_ids"], padding)
    logits = mlm.head(hidden)

real = ~padding  # padded positions hold numbers nobody reads
hidden_gap = (hidden - theirs.hidden_states[-1])[real].abs().max().item()
logits_gap = (logits - theirs.logits)[real].abs().max().item()
print(f"largest difference, hidden states: {hidden_gap:.1e}")
print(f"largest difference, word scores:   {logits_gap:.1e}")
assert hidden_gap < 1e-3 and logits_gap < 1e-2
del reference

largest difference, hidden states: 8.4e-05
largest difference, word scores:   1.1e-04

The differences are rounding: the two implementations add the same numbers in a different order. Our sixty lines are BERT.

5. Filling in the blanks

Pre-training picks 15% of the tokens of every sequence at random, never [CLS] or [SEP], and asks the model to restore them. If every picked token became [MASK], the model would learn something slightly wrong: that only [MASK] positions need predicting, while in fine-tuning [MASK] never appears.

The paper calls this “a mismatch between pre-training and fine-tuning” and softens it with a rule that every BERT-like model has copied since: of the picked tokens, 80% become [MASK], 10% become a random token, and 10% stay as they are.

The model never knows which tokens are real, so it has to keep a good guess for every position.

The 80/10/10 rule. Of all tokens, 15% are picked for prediction. Of those, 80% are replaced with the MASK token, 10% with a random word, and 10% are left unchanged; the model must predict the original word at all of them The 80/10/10 rule. Of all tokens, 15% are picked for prediction. Of those, 80% are replaced with the MASK token, 10% with a random word, and 10% are left unchanged; the model must predict the original word at all of them

def mask_tokens(ids, special, generator, rate=0.15):
    chosen = (torch.rand(ids.shape, generator=generator) < rate) & ~special
    labels = ids.clone()
    labels[~chosen] = -100  # the loss ignores positions that weren't picked
    roll = torch.rand(ids.shape, generator=generator)
    masked = ids.clone()
    masked[chosen & (roll < 0.8)] = tokenizer.mask_token_id
    swap = chosen & (roll >= 0.8) & (roll < 0.9)
    random_ids = torch.randint(len(tokenizer), ids.shape, generator=generator)
    masked[swap] = random_ids[swap]
    return masked, labels


TEXT = ("In October 2019 Google changed the way its search engine reads queries. The old "
        "system treated the little word to as noise and returned results about American "
        "citizens travelling to Brazil, which is the other way around. The new one "
        "understood that the traveller goes from Brazil to the United States.")
encoded = tokenizer(TEXT, return_tensors="pt", return_special_tokens_mask=True)
ids = encoded["input_ids"]
special = encoded["special_tokens_mask"].bool()
generator = torch.Generator().manual_seed(7)
masked, labels = mask_tokens(ids, special, generator)
before = tokenizer.convert_ids_to_tokens(ids[0])
after = tokenizer.convert_ids_to_tokens(masked[0])
token_chips([before, after], labels=["original", "input"], max_width=640,
            strong=lambda row, col, token: bool(labels[0, col] != -100))

Output of the code above

Output of the code above

The highlighted positions are the ones the loss looks at: some became [MASK], and with this many tokens you may also see a word replaced by a random one, or a picked word left alone. Over many sequences the shares settle at what the rule says:

many = ids.repeat(4000, 1)
masked_many, labels_many = mask_tokens(many, special.repeat(4000, 1), generator)
picked = labels_many != -100
print(f"picked for prediction: {picked.float().mean():.1%} of all tokens")
became_mask = (masked_many[picked] == tokenizer.mask_token_id).float().mean()
unchanged = (masked_many[picked] == many[picked]).float().mean()
print(f"  became [MASK]:   {became_mask:.1%}")
print(f"  left unchanged:  {unchanged:.1%}")
print(f"  became another:  {1 - became_mask - unchanged:.1%}")

picked for prediction: 14.4% of all tokens
  became [MASK]:   79.9%
  left unchanged:  10.3%
  became another:  9.9%

A random replacement can occasionally draw the original word back, so “left unchanged” comes out a hair above 10%. Google’s own data generator picked exactly 15% of each sequence rather than tossing a coin per token; the coin is what most libraries do today, and it makes no difference that anyone has measured.

The share picked is a little under 15% because [CLS] and [SEP] are never picked.

The loss is cross-entropy over the vocabulary at the picked positions. On a text BERT has never seen, the first paragraph of this lesson, here is what four days of pre-training bought, against the same network as it was before training: weights drawn at random with a standard deviation of 0.02, the way BERT’s were initialised.

def masked_lm_loss(model, ids, special, repeats=20):
    generator = torch.Generator().manual_seed(0)
    losses = []
    for _ in range(repeats):
        masked, labels = mask_tokens(ids, special, generator)
        segments = torch.zeros_like(masked)
        padding = torch.zeros_like(masked, dtype=torch.bool)
        with torch.no_grad():
            logits = model(masked.to(device), segments.to(device), padding.to(device))
        loss = F.cross_entropy(logits.view(-1, logits.shape[-1]).float(),
                               labels.view(-1).to(device), ignore_index=-100)
        losses.append(loss.item())
    return sum(losses) / len(losses)


def init_like_bert(module):  # how BERT started: small random weights
    if isinstance(module, (nn.Linear, nn.Embedding)):
        nn.init.normal_(module.weight, std=0.02)
    if isinstance(module, nn.Linear) and module.bias is not None:
        nn.init.zeros_(module.bias)
    if isinstance(module, nn.LayerNorm):
        nn.init.ones_(module.weight)
        nn.init.zeros_(module.bias)


untrained = MaskedLM(config)
untrained.apply(init_like_bert)
untrained = untrained.to(device).eval()
print(f"random weights:  loss {masked_lm_loss(untrained, ids, special):.2f}")
print(f"BERT's weights:  loss {masked_lm_loss(mlm, ids, special):.2f}")
print(f"guessing uniformly among {len(tokenizer):,} tokens: {math.log(len(tokenizer)):.2f}")
del untrained

random weights:  loss 10.33
BERT's weights:  loss 3.00
guessing uniformly among 30,522 tokens: 10.33

A loss of \(L\) means the model gives the right word a probability of \(e^{-L}\) on average, in the geometric sense. The untrained network is no better than picking from the dictionary blindfolded.

Four days of pre-training took the loss from 10.33 to 3.00: BERT gives the hidden word a sizeable probability most of the time, on a text written years after it was trained.

That makes BERT a strange kind of oracle: ask it to fill a blank, and you see what it absorbed from books and Wikipedia. Petroni and colleagues used exactly this to show that BERT holds “relational knowledge competitive with traditional NLP methods”[6]. Our own probes:

def fill_mask(text, k=3):
    batch = tokenizer(text, return_tensors="pt").to(device)
    padding = batch["attention_mask"] == 0
    with torch.no_grad():
        logits = mlm(batch["input_ids"], batch["token_type_ids"], padding)
    where = (batch["input_ids"][0] == tokenizer.mask_token_id).nonzero()[0, 0]
    probs = logits[0, where].float().softmax(dim=-1)
    top = probs.topk(k)
    words = tokenizer.convert_ids_to_tokens(top.indices.tolist())
    return ", ".join(f"{w} {p:.0%}" for w, p in zip(words, top.values.tolist()))


PROBES = [
    ("a fact", "The capital of Portugal is [MASK]."),
    ("a fact", "Dante was born in [MASK]."),
    ("common sense", "You cut bread with a [MASK]."),
    ("grammar", "The keys to the old cabinet [MASK] on the kitchen table."),
    ("grammar", "The key to the old cabinets [MASK] on the kitchen table."),
    ("negation", "A robin is a [MASK]."),
    ("negation", "A robin is not a [MASK]."),
    ("negation", "Birds can [MASK]."),
    ("negation", "Birds cannot [MASK]."),
]
table(pd.DataFrame({"probe": [kind for kind, _ in PROBES],
                    "sentence": [text for _, text in PROBES],
                    "BERT's top three": [fill_mask(text) for _, text in PROBES]}))
probesentenceBERT's top three
a factThe capital of Portugal is [MASK].lisbon 73%, porto 20%, madeira 0%
a factDante was born in [MASK].rome 5%, paris 4%, london 2%
common senseYou cut bread with a [MASK].knife 72%, stick 8%, fork 4%
grammarThe keys to the old cabinet [MASK] on the kitchen table.were 52%, sat 17%, lay 12%
grammarThe key to the old cabinets [MASK] on the kitchen table.was 41%, lay 18%, sat 18%
negationA robin is a [MASK].robin 38%, bird 32%, pigeon 1%
negationA robin is not a [MASK].robin 57%, bird 10%, pigeon 1%
negationBirds can [MASK].fly 19%, sing 9%, breed 8%
negationBirds cannot [MASK].fly 37%, sing 10%, escape 3%

The knowledge is patchy in a revealing way. BERT is sure about Lisbon and the knife, and has no idea where Dante was born: a fact that is stated less often in so many words.

The grammar is right, and it is the impressive row. The verb agrees with keys or key, the subject, and not with cabinet or cabinets, the noun right next to it, so the model must know which word the verb belongs to.

Negation is where it breaks. To BERT a robin is above all a robin, and it is more sure that birds cannot fly than that they can.

Kassner and Schütze found this across models and put it in the title of their paper, “Birds Can Talk, But Cannot Fly”: pre-trained models “do not distinguish between negated (‘Birds cannot [MASK]’) and non-negated (‘Birds can [MASK]’) cloze questions”[7].

Filling a blank rewards knowing which words go together, and birds and fly go together whether or not there is a not between them.

The same oracle shows what else came with the books. Give BERT a sentence about a person with a job and ask for the pronoun:

JOBS = ["nurse", "receptionist", "secretary", "librarian", "hairdresser", "teacher",
        "cashier", "baker", "accountant", "doctor", "lawyer", "scientist", "programmer",
        "engineer", "pilot", "mechanic", "carpenter", "surgeon", "firefighter", "farmer"]
he, she = tokenizer.convert_tokens_to_ids(["he", "she"])
shares = {}
for job in JOBS:
    sentence = f"the {job} said that [MASK] would be late."
    batch = tokenizer(sentence, return_tensors="pt").to(device)
    padding = batch["attention_mask"] == 0
    with torch.no_grad():
        logits = mlm(batch["input_ids"], batch["token_type_ids"], padding)
    where = (batch["input_ids"][0] == tokenizer.mask_token_id).nonzero()[0, 0]
    probs = logits[0, where].float().softmax(dim=-1)
    shares[job] = (probs[she] / (probs[he] + probs[she])).item()

order = sorted(shares, key=shares.get)
fig, ax = plt.subplots(figsize=(6.4, 5.2))
ax.barh(order, [shares[j] for j in order], color=COLORS["att"])
ax.axvline(0.5, color=COLORS["muted"], linewidth=1, linestyle="--")
ax.set_xlim(0, 1)
ax.set_title("“The ___ said that [MASK] would be late”: she or he?")
ax.set_xlabel("P(she) / (P(she) + P(he))")
ax.grid(axis="y", visible=False)
plt.show()

“The ___ said that [MASK] would be late”: she or he?

Nothing in the sentence says anything about the person’s gender, and the dashed line is where a model with no opinion would put every job. BERT has opinions.

Receptionists, nurses and hairdressers are she; engineers, pilots and programmers are he; and most jobs, the doctor and the teacher among them, lean he. These are the statistics of its training text.

Kurita and colleagues measured them this way, with templates and the masked word’s probabilities, and showed that such associations carry into the models fine-tuned from BERT [8]. It is worth remembering when a BERT-based classifier screens CVs.

6. What the heads look at

BERT-base has 144 attention heads, 12 in each of 12 layers, and nobody told any of them what to do.

Clark and colleagues looked at all of them and found heads attending “to delimiter tokens, specific positional offsets, or broadly attending over the whole sentence”, and some that “correspond well to linguistic notions of syntax and coreference”[9].

Our model hands back its attention weights, so we can measure the same thing on a few hundred sentences of Wikipedia: for every head, how much of its attention goes to [CLS], to [SEP], to the previous token and to the next one, averaged over all real tokens.

wiki = load_dataset("Salesforce/wikitext", "wikitext-2-raw-v1", split="test")
sentences = [line.strip() for line in wiki["text"] if len(line.strip()) > 200]
sentences = sentences[:32 if SMOKE else 320]

layers, heads = config["num_hidden_layers"], config["num_attention_heads"]
TARGETS = ["[CLS]", "[SEP]", "previous", "next"]
totals = {name: torch.zeros(layers, heads) for name in TARGETS}
count = 0
for start in range(0, len(sentences), 32):
    batch = tokenizer(sentences[start:start + 32], padding=True, truncation=True,
                      max_length=128, return_tensors="pt").to(device)
    padding = batch["attention_mask"] == 0
    with torch.no_grad():
        _, attention = mlm.bert(batch["input_ids"], batch["token_type_ids"], padding,
                                keep_weights=True)
    weights = torch.stack(attention, dim=1).float()  # batch, layer, head, query, key
    real = (~padding).float()
    is_sep = (batch["input_ids"] == tokenizer.sep_token_id).float()
    to_cls = weights[..., 0]
    to_sep = (weights * is_sep[:, None, None, None, :]).sum(-1)
    to_prev = torch.diagonal(weights, offset=-1, dim1=-2, dim2=-1)
    to_next = torch.diagonal(weights, offset=1, dim1=-2, dim2=-1)
    totals["[CLS]"] += (to_cls * real[:, None, None, :]).sum((0, 3)).cpu()
    totals["[SEP]"] += (to_sep * real[:, None, None, :]).sum((0, 3)).cpu()
    totals["previous"] += (to_prev * real[:, None, None, 1:]).sum((0, 3)).cpu()
    totals["next"] += (to_next * real[:, None, None, :-1]).sum((0, 3)).cpu()
    count += real.sum().item()
shares_by_head = {name: (value / count).numpy() for name, value in totals.items()}

fig, axes = plt.subplots(1, 4, figsize=(12, 3.6))
for ax, (name, values) in zip(axes, shares_by_head.items()):
    heatmap(values, rows=range(1, layers + 1), cols=range(1, heads + 1), ax=ax,
            vmin=0, vmax=1, cell=0.22, title=f"to {name}",
            xlabel="head", ylabel="layer" if ax is axes[0] else None)
plt.tight_layout()
plt.show()

Output of the code above

The layers divide the work. In layers 2 to 4 attention gathers on [CLS]; in layers 5 to 10 almost every head sends much of it to [SEP], a token that carries no information about any particular word; the last two layers spread out again.

Clark and colleagues suggested that a head with nothing useful to do in a given context parks its attention on [SEP], a no-op: softmax has to send the weight somewhere.

The first layers also have heads that look almost only one token back or one token ahead, which is how a model with no built-in notion of order puts neighbouring words together. Here is the most devoted previous-token head on one sentence:

layer, head = np.unravel_index(shares_by_head["previous"].argmax(), (layers, heads))
example = "The traveller from Brazil needs a visa to enter the United States."
batch = tokenizer(example, return_tensors="pt").to(device)
with torch.no_grad():
    _, attention = mlm.bert(batch["input_ids"], batch["token_type_ids"],
                            batch["attention_mask"] == 0, keep_weights=True)
words = tokenizer.convert_ids_to_tokens(batch["input_ids"][0])
heatmap(attention[layer][0, head].float().cpu().numpy(), rows=words, cols=words, cell=0.28,
        title=f"Layer {layer + 1}, head {head + 1}: where each word looks",
        xlabel="attends to", ylabel="word")
plt.show()
print(f"on average this head puts {shares_by_head['previous'][layer, head]:.0%} "
      "of its attention on the previous token")

Output of the code above

on average this head puts 71% of its attention on the previous token

Trap

Heads like this one are why reading “attention maps” as explanations is risky. Most of what a head does is plumbing: moving information to where the next layer can use it. The linguistically interesting heads exist, but they are a minority, and a head that looks at a word doesn’t mean the model uses what it found.

7. A job for BERT: personal data in financial documents

Banks, insurers and their software vendors move a lot of text around: support tickets, loan applications, SWIFT messages, contracts. Before such text goes into a log, a search index, a test environment or a prompt for some external model, the personal data in it has to be found and masked: names, addresses, account numbers, phone numbers.

Regular expressions find the things with a fixed shape. Names and addresses have no fixed shape, and that is a classic job for an encoder: named entity recognition, a label for every token.

Real documents full of real personal data are, for good reasons, not something anyone publishes. Gretel’s dataset is synthetic: about 56,000 documents of 60 kinds, in seven languages, written by a language model and annotated with the character spans of 29 kinds of personal data [10]. We take the English ones.

pii = load_dataset("gretelai/synthetic_pii_finance_multilingual")
english = pii.filter(lambda row: row["language"] == "English")
train_rows, test_rows = english["train"], english["test"]
if SMOKE:
    train_rows, test_rows = train_rows.select(range(96)), test_rows.select(range(24))
print(f"{len(train_rows):,} documents to train on, {len(test_rows):,} to test on")
print(f"{len(set(train_rows['document_type']))} kinds, e.g.:",
      ", ".join(sorted(set(train_rows["document_type"]))[:8]))

25,948 documents to train on, 2,962 to test on
60 kinds, e.g.: Annual Report, Audit Report, BAI Format, Bank Statement, Bill of Lading, Business Plan, CSV, Compliance Certificate

The 29 kinds of the annotation are more than we need to see. We fold them into twelve groups, which is roughly how a masking policy would treat them: a first name and a full name are both a person, an IBAN and a routing number are both bank details.

GROUPS = {
    "name": "PERSON", "first_name": "PERSON", "last_name": "PERSON",
    "company": "COMPANY",
    "street_address": "ADDRESS",
    "local_latlng": "LOCATION",
    "date": "DATE", "date_of_birth": "DATE", "date_time": "DATE", "time": "DATE",
    "email": "EMAIL",
    "phone_number": "PHONE",
    "iban": "BANK", "bban": "BANK", "swift_bic_code": "BANK", "bank_routing_number": "BANK",
    "credit_card_number": "CARD", "credit_card_security_code": "CARD",
    "ssn": "GOV_ID", "passport_number": "GOV_ID", "driver_license_number": "GOV_ID",
    "customer_id": "ACCOUNT", "employee_id": "ACCOUNT", "account_pin": "ACCOUNT",
    "ipv4": "DIGITAL", "ipv6": "DIGITAL", "user_name": "DIGITAL", "password": "DIGITAL",
    "api_key": "DIGITAL",
}
KINDS = sorted(set(GROUPS.values()))


def gold_spans(row):
    spans = []
    for span in json.loads(row["pii_spans"]):
        spans.append((span["start"], span["end"], GROUPS[span["label"]]))
    return spans


counts = {kind: 0 for kind in KINDS}
for row in train_rows:
    for _, _, kind in gold_spans(row):
        counts[kind] += 1
order = sorted(KINDS, key=counts.get)
fig, ax = plt.subplots(figsize=(6.4, 3.6))
ax.barh(order, [counts[k] for k in order], color=COLORS["data"])
ax.set_xscale("log")
ax.set_title("Annotated spans in the training documents, by group")
ax.set_xlabel("spans (log scale)")
ax.grid(axis="y", visible=False)
plt.show()

Annotated spans in the training documents, by group

Dates, people and companies are everywhere; GPS coordinates and card numbers are rare. A model will learn the common groups easily and the rare ones from a few hundred examples each. Here is one document with its annotation:

SHADES = {"PERSON": "#D23F57", "COMPANY": "#6A55D8", "ADDRESS": "#0F8C7E",
          "DATE": "#C27A0E"}


def highlight(text, spans):
    pieces, last = [], 0
    for start, end, kind in sorted(spans):
        color = SHADES.get(kind, "#5B6575")
        pieces.append(text[last:start].replace("<", "&lt;"))
        pieces.append(f'<span style="background:{color}2e;border-bottom:2px solid {color};'
                      f'padding:0 2px">{text[start:end]}<sub style="color:{color};'
                      f'font-size:0.7em"> {kind}</sub></span>')
        last = end
    pieces.append(text[last:].replace("<", "&lt;"))
    body = "".join(pieces)
    style = "white-space:pre-wrap;font-size:0.85em;line-height:1.55"
    return HTML(f'<div style="{style}">{body}</div>')


sample = train_rows[3]
print(sample["document_type"])
sample_spans = [s for s in gold_spans(sample) if s[1] <= 900]
display(highlight(sample["generated_text"][:900], sample_spans))

Health Insurance Claim Form
Patient Information * Full Name: Eufemia Gordillo-Pagès PERSON * Date of Birth: MM/DD/YYYY * Address: 8055 Tanya Estates ADDRESS * Policy Number: RUIA78128695550594 BANK Claim Details * Date of Service: MM/DD/YYYY * Provider Name: [Treating Dentist's Name] * Provider Address: [Treating Dentist's Address] ADDRESS * Procedure Code: [Dental Procedure Code] * Procedure Description: [Brief Description of Dental Procedure] * Cost of Service: $[Amount] Notes: * This claim is for a dental check-up and cleaning. * The treating dentist identified no issues during the examination. * Regular dental check-ups are essential for maintaining good oral health. By signing below, I confirm that the information provided is true and accurate to the best of my knowledge. Patient Signature: \_\_\_\_\_\_\_\_\_\_\_\_\_\_\_\_\_\_\_ Date: \_\_\_\_\_\_\_\_\_ Please attach any relevant receipts or invoices to this claim form. Mail the

Synthetic data has its own quirks, and this document shows two. The template line “Provider Address: [Treating Dentist’s Address]” is a placeholder the generator forgot to fill, and the annotator, also a model, marked it as an address. And a policy number is labelled as a bank account number.

A model trained on such labels will learn some of their mistakes too; when we measure it, a small share of its “errors” will be the annotation’s.

The bigger problem is length. BERT reads at most 512 tokens, two of them [CLS] and [SEP]:

lengths = [len(cased_tokenizer(text, add_special_tokens=False)["input_ids"])
           for text in test_rows["generated_text"]]
longer = sum(n > 510 for n in lengths) / len(lengths)
fig, ax = plt.subplots(figsize=(6.4, 3.2))
ax.hist(lengths, bins=40, color=COLORS["data"])
ax.axvline(510, color=COLORS["dec"], linewidth=1.5)
ax.text(520, ax.get_ylim()[1] * 0.85, "BERT's limit", color=COLORS["dec"])
ax.set_title("Length of the test documents in BERT's tokens")
ax.set_xlabel("tokens")
plt.show()
print(f"{longer:.0%} of the documents don't fit in one input")

Length of the test documents in BERT

14% of the documents don't fit in one input

Truncating the 14% of documents that don’t fit would silently skip whatever personal data sits past the limit, which for a masking tool is the worst kind of error.

The standard fix is a sliding window: cut the document into overlapping pieces of 510 tokens, run each, and for every token keep the prediction from the window in which it sits furthest from an edge, where it had the most context on both sides.

A document of 1,100 tokens cut into three overlapping windows of 510 tokens with a stride of 384. Each token's label is taken from the window where it is furthest from the window's edge, so every token is judged with context on both sides A document of 1,100 tokens cut into three overlapping windows of 510 tokens with a stride of 384. Each token's label is taken from the window where it is furthest from the window's edge, so every token is judged with context on both sides

WINDOW, STRIDE = (126, 96) if SMOKE else (510, 384)


def windows(n_tokens):
    if n_tokens <= WINDOW:
        return [(0, n_tokens)]
    starts = list(range(0, n_tokens - WINDOW + 1, STRIDE))
    if starts[-1] + WINDOW < n_tokens:
        starts.append(n_tokens - WINDOW)
    return [(start, start + WINDOW) for start in starts]


print(windows(1100))

[(0, 510), (384, 894), (590, 1100)]

Labels come as character spans, and the model needs one label per token. We use the BIO scheme: B-PERSON on the first token of a name, I-PERSON on the tokens that continue it, O on everything else. The tokenizer’s offset mapping says which characters each token covers:

LABELS = ["O"] + [f"{prefix}-{kind}" for kind in KINDS for prefix in ("B", "I")]
LABEL_ID = {label: i for i, label in enumerate(LABELS)}


def encode(row):
    text = row["generated_text"]
    encoded = cased_tokenizer(text, add_special_tokens=False, return_offsets_mapping=True)
    offsets = encoded["offset_mapping"]
    tags = ["O"] * len(offsets)
    for start, end, kind in gold_spans(row):
        first = True
        for i, (a, b) in enumerate(offsets):
            if a >= start and b <= end and b > a:
                tags[i] = ("B-" if first else "I-") + kind
                first = False
    return encoded["input_ids"], offsets, tags


ids, offsets, tags = encode(sample)
shown = [(cased_tokenizer.convert_ids_to_tokens(i), t) for i, t in zip(ids[:40], tags[:40])]
token_chips([[token for token, _ in shown]], max_width=640,
            strong=lambda row, col, token: shown[col][1] != "O")
print(" ".join(tag for _, tag in shown))

O O O O O O O B-PERSON I-PERSON I-PERSON I-PERSON I-PERSON I-PERSON I-PERSON I-PERSON I-PERSON I-PERSON O O O O O O O O O O O O O O O O O O B-ADDRESS I-ADDRESS I-ADDRESS I-ADDRESS O

Output of the code above

Each window becomes a training example: [CLS], the window’s tokens, [SEP], and labels of \(-100\) on the two special tokens so that the loss ignores them.

def examples(rows):
    out = []
    for row in rows:
        ids, _, tags = encode(row)
        for start, end in windows(len(ids)):
            window_ids = [cased_tokenizer.cls_token_id] + ids[start:end]
            window_ids += [cased_tokenizer.sep_token_id]
            window_labels = [-100] + [LABEL_ID[t] for t in tags[start:end]] + [-100]
            out.append((window_ids, window_labels))
    return out

8. Fine-tuning, and what pre-training is worth

The tagger is our Bert with one linear layer on top: 768 numbers in, one score per label out, at every position. All of BERT’s weights train along with the new layer, which is what “fine-tuning” meant in the paper.

class Tagger(nn.Module):
    def __init__(self, bert, n_labels):
        super().__init__()
        self.bert = bert
        self.dropout = nn.Dropout(0.1)
        self.classify = nn.Linear(bert.embeddings.words.embedding_dim, n_labels)

    def forward(self, ids, padding):
        h, _ = self.bert(ids, torch.zeros_like(ids), padding)
        return self.classify(self.dropout(h))


def new_tagger(pretrained):  # random weights are initialised as in section 5
    masked_lm, _ = load_masked_lm("google-bert/bert-base-cased")
    tagger = Tagger(masked_lm.bert, len(LABELS))
    if not pretrained:
        tagger.apply(init_like_bert)
    return tagger.to(device)

Training sorts the windows by length and batches neighbours together, so a batch of short windows isn’t padded to

  1. The rest is the recipe from the paper’s fine-tuning appendix: AdamW, a learning rate warmed up over the first tenth of the steps and then decayed linearly, a few epochs. Mixed precision keeps it fast on a T4.
def batches(data, size, shuffle):
    order = sorted(range(len(data)), key=lambda i: len(data[i][0]))
    groups = [order[i:i + size] for i in range(0, len(order), size)]
    if shuffle:
        random.shuffle(groups)
    for group in groups:
        width = max(len(data[i][0]) for i in group)
        ids = torch.full((len(group), width), cased_tokenizer.pad_token_id)
        labels = torch.full((len(group), width), -100)
        for row, i in enumerate(group):
            n = len(data[i][0])
            ids[row, :n] = torch.tensor(data[i][0])
            labels[row, :n] = torch.tensor(data[i][1])
        yield ids.to(device), labels.to(device)


def train(tagger, data, epochs=2, lr=5e-5, batch_size=16, steps=None):
    if steps:  # a fixed number of updates, however much data there is
        epochs = math.ceil(steps / math.ceil(len(data) / batch_size))
    steps = epochs * math.ceil(len(data) / batch_size)
    warmup = max(1, steps // 10)
    optimizer = torch.optim.AdamW(tagger.parameters(), lr=lr, weight_decay=0.01)
    def rate(step):  # linear warm-up, then linear decay to zero
        return min((step + 1) / warmup, max(0.0, (steps - step) / max(1, steps - warmup)))

    schedule = torch.optim.lr_scheduler.LambdaLR(optimizer, rate)
    scaler = torch.amp.GradScaler(enabled=device == "cuda")
    tagger.train()
    for _ in range(epochs):
        for ids, labels in batches(data, batch_size, shuffle=True):
            with torch.autocast(device, dtype=torch.float16, enabled=device == "cuda"):
                logits = tagger(ids, ids == cased_tokenizer.pad_token_id)
                flat = logits.view(-1, len(LABELS)).float()
                loss = F.cross_entropy(flat, labels.view(-1), ignore_index=-100)
            optimizer.zero_grad(set_to_none=True)
            scaler.scale(loss).backward()
            scaler.unscale_(optimizer)
            nn.utils.clip_grad_norm_(tagger.parameters(), 1.0)
            scaler.step(optimizer)
            scaler.update()
            schedule.step()
    return tagger.eval()

Prediction runs the windows of a document, keeps each token’s best-placed prediction, and turns runs of B- and I- labels back into character spans. Evaluation is strict: a predicted span counts only if it has the right group and exactly the right boundaries, on the token grid.

def predict(tagger, row):
    ids, offsets, _ = encode(row)
    best, depth = ["O"] * len(ids), [-1] * len(ids)
    for start, end in windows(len(ids)):
        window = [cased_tokenizer.cls_token_id] + ids[start:end]
        window = torch.tensor([window + [cased_tokenizer.sep_token_id]], device=device)
        no_padding = torch.zeros_like(window, dtype=torch.bool)
        half = torch.autocast(device, dtype=torch.float16, enabled=device == "cuda")
        with torch.no_grad(), half:
            logits = tagger(window, no_padding)
        labels = logits[0, 1:-1].argmax(-1).tolist()
        for k, label in enumerate(labels):
            i = start + k
            room = min(k, end - start - 1 - k)  # distance to the nearer edge of the window
            if room > depth[i]:
                best[i], depth[i] = LABELS[label], room
    return to_spans(best, offsets)


def to_spans(tags, offsets):
    spans, current = [], None
    for (a, b), tag in zip(offsets, tags):
        kind = tag[2:] if tag != "O" else None
        if kind and (tag.startswith("B-") or current is None or current[2] != kind):
            if current:
                spans.append(tuple(current))
            current = [a, b, kind]
        elif kind:
            current[1] = b
        elif current:
            spans.append(tuple(current))
            current = None
    if current:
        spans.append(tuple(current))
    return spans


def aligned_gold(row):
    _, offsets, tags = encode(row)
    return to_spans(tags, offsets)


def scores(predicted, gold):
    rows = []
    for kind in KINDS + ["all"]:
        hits = found = wanted = 0
        for p, g in zip(predicted, gold):
            p = {s for s in p if kind == "all" or s[2] == kind}
            g = {s for s in g if kind == "all" or s[2] == kind}
            hits, found, wanted = hits + len(p & g), found + len(p), wanted + len(g)
        precision = hits / found if found else float("nan")
        recall = hits / wanted if wanted else float("nan")
        f1 = 2 * precision * recall / (precision + recall) if hits else 0.0
        rows.append({"group": kind, "spans": wanted, "precision": precision,
                     "recall": recall, "F1": f1})
    return pd.DataFrame(rows)

First the experiment that explains why BERT mattered. The same network is fine-tuned twice on the same documents: once starting from Google’s weights, once from random ones, each with the same 300 updates of the weights. We repeat it with more and more training documents.

random.seed(SEED)
test_subset = test_rows.select(range(min(len(test_rows), 400)))
test_gold = [aligned_gold(row) for row in test_subset]
SIZES = [16] if SMOKE else [100, 400, 1600]  # CI runs on a slow CPU
curve = []
for size in SIZES:
    data = examples(train_rows.select(range(size)))
    for pretrained in (True, False):
        t0 = time.time()
        torch.manual_seed(SEED)
        tagger = train(new_tagger(pretrained), data, steps=10 if SMOKE else 300)
        predicted = [predict(tagger, row) for row in test_subset]
        f1 = scores(predicted, test_gold).set_index("group").loc["all", "F1"]
        start = "BERT's weights" if pretrained else "random"
        curve.append({"documents": size, "start": start, "F1": f1,
                      "minutes": (time.time() - t0) / 60})
        del tagger
        if device == "cuda":
            torch.cuda.empty_cache()
table(pd.DataFrame(curve))
documentsstartF1minutes
100BERT's weights0.532.33
100random0.052.42
400BERT's weights0.592.76
400random0.072.82
1,600BERT's weights0.643.12
1,600random0.083.16
curve_frame = pd.DataFrame(curve)
fig, ax = plt.subplots(figsize=(6.4, 3.6))
for start, color in [("BERT's weights", COLORS["enc"]), ("random", COLORS["dec"])]:
    part = curve_frame[curve_frame["start"] == start]
    ax.plot(part["documents"], part["F1"], marker="o", color=color, label=start)
ax.set_xscale("log")
ax.set_xticks(SIZES, [str(s) for s in SIZES])
ax.set_ylim(0, 1)
ax.set_title("Same network, same data: starting from pre-trained weights")
ax.set_xlabel("training documents")
ax.set_ylabel("entity F1 on the test documents")
ax.legend(frameon=False)
plt.show()

Same network, same data: starting from pre-trained weights

This is the chart that changed NLP in 2018. With a hundred documents the pre-trained network already gets about half of the spans exactly right.

The random one gets almost none, and sixteen times more documents barely move it: it spends its 300 updates learning that most tokens are O, while the pre-trained one starts from a model of English and only has to learn what a name or an IBAN looks like in these documents.

With far more updates and data the random network would get further; pre-training paid for those in advance.

Key point

The four days of Google’s TPUs are not in the architecture, since the two networks are identical. They are in the weights.

Now the real model, on all the training documents:

train_data = examples(train_rows)
print(f"{len(train_data):,} windows from {len(train_rows):,} documents")
t0 = time.time()
torch.manual_seed(SEED)
tagger = train(new_tagger(pretrained=True), train_data, epochs=1 if SMOKE else 2)
print(f"fine-tuned in {(time.time() - t0) / 60:.1f} min")

30,142 windows from 25,948 documents
fine-tuned in 30.7 min

9. BERT against regular expressions

The fair opponent is what a careful engineer would write first: regular expressions for everything that has a shape. Emails, phone numbers, IBANs, card numbers checked with the Luhn algorithm, dates in the common formats, IP addresses, social security numbers. Names, companies and addresses get no regex, because there is no honest one to write.

MONTHS = "Jan|Feb|Mar|Apr|May|Jun|Jul|Aug|Sep|Oct|Nov|Dec"
PATTERNS = {
    "EMAIL": r"[\w.+-]+@[\w-]+(?:\.[\w-]+)+",
    "PHONE": r"(?<![\w-])(?:\+\d{1,3}[\s.-]?)?(?:\(\d{2,4}\)[\s.-]?)?\d{3}[\s.-]\d{3,4}"
             r"(?:[\s.-]\d{2,4})?(?![\w-])",
    "BANK": r"\b[A-Z]{2}\d{2}(?:\s?[A-Z0-9]{4}){2,7}(?:\s?[A-Z0-9]{1,4})?\b",
    "CARD": r"\b(?:\d[ -]?){12,18}\d\b",
    "DATE": rf"\b\d{{4}}-\d{{2}}-\d{{2}}\b|\b\d{{1,2}}/\d{{1,2}}/\d{{2,4}}\b"
            rf"|\b(?:{MONTHS})[a-z]*\.? \d{{1,2}},? \d{{4}}\b"
            rf"|\b\d{{1,2}} (?:{MONTHS})[a-z]* \d{{4}}\b"
            rf"|\b\d{{1,2}}:\d{{2}}(?::\d{{2}})?(?:\s?[AP]M)?\b",
    "DIGITAL": r"\b(?:\d{1,3}\.){3}\d{1,3}\b|\b(?:[0-9a-fA-F]{1,4}:){7}[0-9a-fA-F]{1,4}\b",
    "GOV_ID": r"\b\d{3}-\d{2}-\d{4}\b",
}


def luhn(digits):
    total = 0
    for i, d in enumerate(reversed(digits)):
        d = int(d)
        if i % 2 == 1:
            d = d * 2 - 9 if d * 2 > 9 else d * 2
        total += d
    return total % 10 == 0


def regex_spans(row):
    text = row["generated_text"]
    _, offsets, _ = encode(row)
    found = []
    for kind, pattern in PATTERNS.items():
        for match in re.finditer(pattern, text):
            if kind == "CARD" and not luhn(re.sub(r"\D", "", match.group())):
                continue
            inside = [(a, b) for a, b in offsets if a >= match.start() and b <= match.end()]
            if inside:  # the same token grid as the model's predictions
                found.append((inside[0][0], inside[-1][1], kind))
    return found


test_all = test_rows
gold_all = [aligned_gold(row) for row in test_all]
t0 = time.time()
bert_found = [predict(tagger, row) for row in test_all]
seconds = time.time() - t0
regex_found = [regex_spans(row) for row in test_all]
bert_scores = scores(bert_found, gold_all).set_index("group")
regex_scores = scores(regex_found, gold_all).set_index("group")
comparison = pd.DataFrame({"spans": bert_scores["spans"],
                           "BERT F1": bert_scores["F1"],
                           "regex F1": regex_scores["F1"].where(
                               regex_scores.index.isin(list(PATTERNS) + ["all"]))})
display(table(comparison.reset_index().fillna("—")))  # no regex for that group
print(f"BERT tagged {len(test_all):,} documents in {seconds:.0f} s, "
      f"{len(test_all) / seconds:.0f} documents a second")
groupspansBERT F1regex F1
ACCOUNT2810.81—
ADDRESS2,3300.86—
BANK4250.870.31
CARD1750.740.31
COMPANY3,6060.78—
DATE6,6380.820.45
DIGITAL2890.760.69
EMAIL8440.860.93
GOV_ID2250.850.45
LOCATION670.59—
PERSON4,5550.85—
PHONE5320.900.58
all19,9670.820.30
BERT tagged 2,962 documents in 51 s, 58 documents a second
shown_kinds = [k for k in KINDS if bert_scores.loc[k, "spans"] > 0]
x = np.arange(len(shown_kinds))
fig, ax = plt.subplots(figsize=(8.4, 3.6))
ax.bar(x - 0.2, [bert_scores.loc[k, "F1"] for k in shown_kinds], 0.4, color=COLORS["enc"],
       label="fine-tuned BERT")
regex_values = [regex_scores.loc[k, "F1"] if k in PATTERNS else 0 for k in shown_kinds]
ax.bar(x + 0.2, regex_values, 0.4, color=COLORS["data"], label="regular expressions")
ax.set_xticks(x, shown_kinds, rotation=35, ha="right")
ax.set_ylim(0, 1)
ax.set_title("Entity F1 on the test documents, by group")
ax.legend(frameon=False, loc="upper left", bbox_to_anchor=(1, 1))
plt.show()

Entity F1 on the test documents, by group

Look at emails first: there the regular expression beats the model. An email address has a strict shape, and a pattern matches it character by character, while BERT reads tokens and now and then stops a piece short. Everywhere else the model is ahead, by the most on bank details, card numbers, ID numbers and dates.

The regular expressions lose there on both counts: they miss every format nobody wrote a pattern for, a passport number, an account number without a country code, a date spelled out in words, and they flag strings of the right shape that are nobody’s personal data.

And for people, companies and addresses, most of the personal data in these documents, there is no pattern to compare with at all. BERT reads the context: Full Name: before a string makes it a person, Ship to before a line makes it an address.

In practice

The two are combined: patterns for the few strict shapes, where they are exact and easy to audit, and the model for everything else.

Here is the model at work on a test document it has never seen, everything it found replaced by its group:

def redact(text, spans):
    out, last = [], 0
    for start, end, kind in sorted(spans):
        out.append(text[last:start])
        out.append(f"[{kind}]")
        last = end
    out.append(text[last:])
    return "".join(out)


# a test document with something of every common group in its first 1,200 characters,
# and none of the generator's unfilled template fields, such as "[Graduation Date]"
for k in range(len(test_all)):
    early = [s for s in bert_found[k] if s[1] <= 1200]
    filled = "[" not in test_all[k]["generated_text"][:1200]
    if filled and {"PERSON", "COMPANY", "ADDRESS", "DATE"} <= {s[2] for s in early}:
        break
demo, demo_spans = test_all[k], early
print(demo["document_type"])
display(highlight(demo["generated_text"][:1200], demo_spans))
print(redact(demo["generated_text"][:1200], demo_spans))

Credit Application
--- Credit Application - Medical Loan Full Name: John Smith PERSON Date of Birth: 01/01/1980 DATE Address: 123 Maple Street, Anytown, USA ADDRESS Email: john.smith@email.com EMAIL Phone: (123) 456-7890 PHONE Employment Information Employer Name: XYZ Corporation COMPANY Address: 456 Oak Street, Anytown, USA ADDRESS Job Title: Software Engineer Monthly DATE Income: $7,500 Loan Details Loan Amount: $25,000 Loan Purpose: Cover medical expenses for upcoming surgery Medical Provider Information Provider Name: Dr. Jane Doe PERSON Address: 789 Pine Street, Anytown, USA ADDRESS Phone: (987) 654-3210 PHONE Treatment Details Description of Medical Services: - Pre-surgery consultation and testing - Complex knee surgery - Post-surgery rehabilitation Cost Breakdown - Pre-surgery consultation and testing: $3,500 - Surgery: $20,000 - Rehabilitation: $1,500 - Total: $25,000 I, John Smith PERSON, hereby authorize the release of my medical and financial information to the necessary parties for the purpose of processing this loan application. Signature: John Smith PERSON Date: 01/15/2023 DATE ---
---
Credit Application - Medical Loan

Full Name: [PERSON]
Date of Birth: [DATE]
Address: [ADDRESS]
Email: [EMAIL]
Phone: [PHONE]

Employment Information

Employer Name: [COMPANY]
Address: [ADDRESS]
Job Title: Software Engineer
[DATE] Income: $7,500

Loan Details

Loan Amount: $25,000
Loan Purpose: Cover medical expenses for upcoming surgery

Medical Provider Information

Provider Name: [PERSON]
Address: [ADDRESS]
Phone: [PHONE]

Treatment Details

Description of Medical Services:
- Pre-surgery consultation and testing
- Complex knee surgery
- Post-surgery rehabilitation

Cost Breakdown

- Pre-surgery consultation and testing: $3,500
- Surgery: $20,000
- Rehabilitation: $1,500
- Total: $25,000

I, [PERSON], hereby authorize the release of my medical and financial information to the necessary parties for the purpose of processing this loan application.

Signature: [PERSON]
Date: [DATE]

---

One mistake is in plain sight: in “Monthly Income” the model took Monthly for a date, a word it has seen next to dates many times.

Two caveats belong next to any number above. The test documents come from the same generator as the training ones, with the same templates and the same habits, so the scores are an upper bound for real documents, which will be messier. And a share of the “errors” are the annotation’s, as we saw in section 7.

Before a model like this masks real data, it needs a few hundred real documents annotated by people, and a threshold chosen for recall, since a missed account number costs more than a masked word.

10. Where BERT is the right tool today

BERT was overtaken years ago on every leaderboard, and it is still everywhere.

For the jobs it was built for, classifying text and labelling tokens, a fine-tuned encoder of a hundred million parameters is usually as accurate as a large language model asked nicely, and far cheaper: our tagger went through almost three thousand documents in under a minute on a free T4, while a 70-billion-parameter model has more than six hundred times as many weights to run for every token it writes.

BERT reads the whole input at once instead of producing text token by token, its output is a label rather than prose that must be parsed, and it runs on a CPU if it has to.

Where it is the wrong tool: anything that needs generation, anything longer than 512 tokens that can’t be cut into windows, and languages it didn’t see in pre-training (multilingual BERT covers 104 languages, worse than the English one covers English).

Its 2018 training recipe has also aged; the next lesson takes it apart and shows what seven years of improvements did to it, from RoBERTa to ModernBERT.

Every lesson of this course ends with the same card, the architecture’s passport, so that the models can be compared side by side:

BERT's passport. Published October 2018 by Google. Type: encoder. Pre-training objective: masked language modelling and next sentence prediction. Attention: bidirectional, full. Positions: learned, up to 512. Normalisation: LayerNorm after each sublayer. Activation: GELU. Sizes: base 110 million parameters, 12 layers, width 768; large 340 million, 24 layers, width 1024. Best at: classification, tagging, retrieval BERT's passport. Published October 2018 by Google. Type: encoder. Pre-training objective: masked language modelling and next sentence prediction. Attention: bidirectional, full. Positions: learned, up to 512. Normalisation: LayerNorm after each sublayer. Activation: GELU. Sizes: base 110 million parameters, 12 layers, width 768; large 340 million, 24 layers, width 1024. Best at: classification, tagging, retrieval

and with the family tree, which grows by one branch per lesson:

The family of the Transformer as this course covers it, from the 2017 Transformer: encoders (BERT, RoBERTa to ModernBERT, sentence embeddings), decoders (GPT, modern LLMs, mixtures of experts, long context, Mamba hybrids), encoder–decoders (T5 and BART, Whisper) and models beyond text (ViT, CLIP, vision-language models, diffusion transformers). BERT, this lesson, is highlighted The family of the Transformer as this course covers it, from the 2017 Transformer: encoders (BERT, RoBERTa to ModernBERT, sentence embeddings), decoders (GPT, modern LLMs, mixtures of experts, long context, Mamba hybrids), encoder–decoders (T5 and BART, Whisper) and models beyond text (ViT, CLIP, vision-language models, diffusion transformers). BERT, this lesson, is highlighted

References

  1. Nayak, P. (2019). Understanding searches better than ever before. Google, 25 October 2019. ↩
  2. Devlin, J., Chang, M.-W., Lee, K., & Toutanova, K. (2018). BERT: pre-training of deep bidirectional transformers for language understanding. NAACL 2019. ↩
  3. Peters, M. E., Neumann, M., Iyyer, M., Gardner, M., Clark, C., Lee, K., & Zettlemoyer, L. (2018). Deep contextualized word representations. NAACL 2018. ↩
  4. Radford, A., Narasimhan, K., Salimans, T., & Sutskever, I. (2018). Improving language understanding by generative pre-training. OpenAI. ↩
  5. Taylor, W. L. (1953). “Cloze procedure”: a new tool for measuring readability. Journalism Quarterly, 30(4), 415–433. ↩
  6. Petroni, F., Rocktäschel, T., Lewis, P., Bakhtin, A., Wu, Y., Miller, A. H., & Riedel, S. (2019). Language models as knowledge bases? EMNLP 2019. ↩
  7. Kassner, N., & Schütze, H. (2020). Negated and misprimed probes for pretrained language models: birds can talk, but cannot fly. ACL 2020. ↩
  8. Kurita, K., Vyas, N., Pareek, A., Black, A. W., & Tsvetkov, Y. (2019). Measuring bias in contextualized word representations. 1st Workshop on Gender Bias in Natural Language Processing. ↩
  9. Clark, K., Khandelwal, U., Levy, O., & Manning, C. D. (2019). What does BERT look at? An analysis of BERT’s attention. BlackboxNLP 2019. ↩
  10. Gretel (2024). Synthetic PII finance multilingual. Hugging Face dataset, Apache 2.0. ↩
]]> Eva Koroleva Mixture of experts: why the largest LLMs use a few percent of their weights per token 2026-02-06T09:00:00+00:00 2026-11-02T00:00:00+00:00 https://blog.korolevalab.com/blog/2026/mixture-of-experts

In short

A mixture of experts splits each feed-forward network into many smaller ones and lets a router send every token to a few of them, so the newest models hold twenty to thirty times as many parameters as they use for a token.

In my small runs, a mixture with the compute of a small dense model got two thirds of the way to a dense model with as many parameters, once its router was kept from piling the tokens onto a few favourite experts.

Serving is where a mixture is expensive: every expert must sit in memory, and a server’s batch reads nearly all of them at every step.

31%longer per training step than a dense model doing the same arithmetic
94–98%of the tokens one expert takes in three of six layers, without any balancing
87%of DeepSeek-V3's routed experts read in one decoding step for 64 sequences

DeepSeek-V3 has 671 billion parameters and uses 37 billion of them for each token. Kimi K2 has a trillion and uses 32 billion, gpt-oss-120b uses 5.1 billion of its 117, and some names spell the arithmetic out: Qwen3-235B-A22B has 235 billion parameters, 22 billion of them active.

In every layer of these models a small router looks at the token and picks a few feed-forward networks, the experts, out of dozens or hundreds. Only those are computed. The rest of the weights wait in memory for tokens that need them.

The idea is older than the Transformer. In 1991 Jacobs, Jordan, Nowlan and Hinton described systems “composed of many separate networks, each of which learns to handle a subset of the complete set of training cases”[1].

In 2017 Shazeer and colleagues put a layer of up to 131,072 experts between the layers of an LSTM language model[2], and today the largest open models of DeepSeek, Qwen, Moonshot, OpenAI, Meta, Zhipu and MiniMax are all mixtures of experts.

This post builds a mixture-of-experts layer in PyTorch, trains it next to dense models of the same cost on English, Russian and Python code, and looks inside: what the router decides, why it has to be kept from playing favourites, what the experts end up specialising in, and what it takes to serve a model whose weights are mostly idle.

Where a Transformer keeps its parameters

A layer of a Transformer has two parts. Attention moves information between tokens: each token looks at the earlier ones and takes what it needs from them.

The feed-forward network then works on every token separately, with the same weights for all of them: it multiplies the token’s vector by a wide matrix, passes the result through a nonlinearity and multiplies it by a second matrix back to the original width. In modern models it is a SwiGLU network with three matrices:

\[\mathrm{FFN}(x) = W_{\text{down}}\big(\mathrm{silu}(W_{\text{gate}}\, x) \odot W_{\text{up}}\, x\big).\]

Here \(x\) is the token’s vector of width \(d\), \(W_{\text{gate}}\) and \(W_{\text{up}}\) map it to a hidden width \(h\), usually three to five times \(d\), the elementwise product \(\odot\) lets one projection switch the other on and off, and \(W_{\text{down}}\) maps the result back to width \(d\).

Most of a model’s weights are in these networks. Qwen3-32B has 64 layers of width 5,120 with feed-forward networks of width 25,600, and its config gives the exact split:

Parameters of two Qwen3 models by part, from their config files. Qwen3-32B, dense: attention 6.0 billion, feed-forward networks 25.2 billion, embeddings 1.6 billion, all used for every token. Qwen3-235B-A22B, a mixture of experts: attention 6.7 billion, the experts a token uses 14.2 billion, the experts it doesn't use 212.9 billion, embeddings 1.2 billion Parameters of two Qwen3 models by part, from their config files. Qwen3-32B, dense: attention 6.0 billion, feed-forward networks 25.2 billion, embeddings 1.6 billion, all used for every token. Qwen3-235B-A22B, a mixture of experts: attention 6.7 billion, the experts a token uses 14.2 billion, the experts it doesn't use 212.9 billion, embeddings 1.2 billion

The feed-forward networks hold 25.2 billion of its 32.8 billion parameters, 77%; attention has 6.0 billion and the embeddings 1.6. Qwen3-235B-A22B, a mixture of experts of the same family, turns the picture around. Of its 235 billion parameters a token passes through 22 billion, and 213 billion, 91% of the model, are experts that this token does not use.

Every weight a token passes through costs two operations, a multiplication and an addition, so a dense model of 32 billion parameters spends about 64 billion operations on every token it reads or writes, and most of them in the feed-forward networks.

What do these networks do with so many weights? Geva and colleagues read the two matrices as a memory.

Each row of the first matrix is a key that responds to some pattern in the input, each column of the second is a value added to the output when its key responds, and the network’s output is the sum of the values whose keys fired.

In their words, feed-forward layers “operate as key-value memories, where each key correlates with textual patterns in the training examples”[3].

A memory of 25,600 patterns of which a given token needs a few, read in full for every token, is an obvious candidate for economy. The idea of the mixture of experts is to cut it into many smaller networks and let every token use only some of them.

A layer of experts

A mixture-of-experts layer replaces the one feed-forward network with \(E\) smaller ones, the experts, and adds a router: a matrix \(W_r\) that gives every expert a score for the current token. A softmax turns the scores into probabilities, the token goes to the \(k\) experts with the highest ones, and its output is their outputs added with weights:

\[p = \mathrm{softmax}(W_r\, x), \qquad \mathcal{T} = \text{the } k \text{ largest of } p_1, \dots, p_E, \qquad y = \sum_{i \in \mathcal{T}} \frac{p_i}{\sum_{j \in \mathcal{T}} p_j}\, E_i(x).\]

Term by term: \(W_r\) is an \(E \times d\) matrix, so the router is a single linear layer, tiny next to the experts. \(p_i\) is how much the router trusts expert \(i\) with this token.

\(\mathcal{T}\) is the set of the \(k\) chosen experts; the others are not computed at all, which is the whole point. The weights are the chosen probabilities divided by their sum, so they add up to one, and each \(E_i\) is an ordinary SwiGLU network, only narrower than the dense one.

A dense feed-forward layer next to a mixture of experts. In the dense layer every token passes through all 768 hidden units. In the mixture, a router scores 8 experts of 384 hidden units each for the token, the token goes to the 2 best, and its output is their outputs added with the router's weights. The token passes through 768 hidden units in both, but the mixture holds 3072 A dense feed-forward layer next to a mixture of experts. In the dense layer every token passes through all 768 hidden units. In the mixture, a router scores 8 experts of 384 hidden units each for the token, the token goes to the 2 best, and its output is their outputs added with the router's weights. The token passes through 768 hidden units in both, but the mixture holds 3072

The models trained below have experts of width 384 and send every token to 2 of 8, so a token passes through 768 hidden units, as in a dense layer of width 768, while the layer holds 3,072.

Here is what the router of one of them, after training, does with the Russian word “президент” in a sentence it has not seen:

The router of layer 3 of the trained model with 8 experts, deciding where the token 'президент' of a held-out Russian sentence goes. Its probabilities for the eight experts are E1 0.07, E2 0.01, E3 0.04, E4 0.10, E5 0.36, E6 0.32, E7 0.04, E8 0.06. The two largest, E5 and E6, are kept and divided by their sum, giving weights 0.53 and 0.47; the token's output is the two experts' outputs added with these weights The router of layer 3 of the trained model with 8 experts, deciding where the token 'президент' of a held-out Russian sentence goes. Its probabilities for the eight experts are E1 0.07, E2 0.01, E3 0.04, E4 0.10, E5 0.36, E6 0.32, E7 0.04, E8 0.06. The two largest, E5 and E6, are kept and divided by their sum, giving weights 0.53 and 0.47; the token's output is the two experts' outputs added with these weights

The router hesitates over this token. Experts 5 and 6 get 0.36 and 0.32 of the probability, and the other six share the remaining third, expert 4 most of it with 0.10.

The two largest probabilities are kept and divided by their sum, 0.68, which turns them into weights 0.53 and 0.47, so the token’s output is roughly half of what expert 5 computes and half of what expert 6 computes. Experts 1, 2, 3, 4, 7 and 8 never see the token at all.

The router learns with the rest of the model. Its gradient comes through the weights \(p_i / \sum_j p_j\): if expert 3’s output would have lowered the loss more than expert 7’s, the gradient raises \(p_3\) relative to \(p_7\), and the next similar token is more likely to go to expert 3.

Nothing tells the router what an expert should be good at. Specialisation, if any, appears because sending similar tokens to the same expert makes that expert better at them.

In code the layer is a loop over experts, each taking the tokens that chose it:

def moe(x, router, experts, k):
    probs = (x @ router).softmax(dim=-1)              # (tokens, experts)
    top_probs, top_experts = probs.topk(k, dim=-1)    # the k best experts of every token
    weights = top_probs / top_probs.sum(dim=-1, keepdim=True)
    out = torch.zeros_like(x)
    for e, expert in enumerate(experts):
        rows, slot = torch.nonzero(top_experts == e, as_tuple=True)   # tokens that chose e
        out[rows] += weights[rows, slot, None] * expert(x[rows])
    return out

Whether this loop is slow depends on how many tokens each expert gets. Every pass launches a few matrix products, and a GPU is efficient only when a product is large enough to keep all of its cores busy.

In training, with thousands of tokens in a batch and a handful of experts, each expert gets plenty of them. With hundreds of experts, or in decoding, where a step has one token per sequence, each expert gets a few rows, and the loop spends its time launching kernels and gathering rows instead of multiplying.

The usual remedy is to sort the token-expert pairs by expert, so that every expert’s tokens lie next to each other, and to run all the experts at once:

How a mixture-of-experts layer processes a batch. The first 8 tokens of a held-out English sentence each go to two of the eight experts of layer 3: 'At' to E6 and E5; 'the' to E3 and E1; 'beginning' to E1 and E4; 'of' to E5 and E3; 'the' to E1 and E3; '21' to E1 and E6; 'st' to E7 and E6; 'century' to E6 and E7. Sorted by expert, they form a queue for each expert: E1: 4; E2: 0; E3: 3; E4: 1; E5: 2; E6: 4; E7: 2; E8: 0. The queues are padded to the same length and every expert processes its queue in one batched product How a mixture-of-experts layer processes a batch. The first 8 tokens of a held-out English sentence each go to two of the eight experts of layer 3: 'At' to E6 and E5; 'the' to E3 and E1; 'beginning' to E1 and E4; 'of' to E5 and E3; 'the' to E1 and E3; '21' to E1 and E6; 'st' to E7 and E6; 'century' to E6 and E7. Sorted by expert, they form a queue for each expert: E1: 4; E2: 0; E3: 3; E4: 1; E5: 2; E6: 4; E7: 2; E8: 0. The queues are padded to the same length and every expert processes its queue in one batched product

How my version batches the experts, and what production kernels do instead

My version pads every expert’s queue to the length of the longest, rounded up to a multiple of 256 so that the GPU sees few different shapes, and multiplies all of them by the stacked expert weights in one batched product. It agrees with a reference that runs every expert on every token and multiplies the unchosen ones by zero, to rounding error. Production kernels avoid even the padding: MegaBlocks writes the layer as a block-sparse matrix product[4], and DeepSeek’s DeepGEMM library has grouped products that take the experts’ tokens concatenated one after another, each expert’s segment padded only to the kernel’s block size.

Same compute, more parameters

What does a mixture buy? To see it on a small scale, I trained four models on a corpus of three kinds of text: English Wikipedia, Russian Wikipedia and the Python standard library, 8 million tokens of each, with a tokenizer of 8,192 tokens trained on the same mix (code).

All four have 6 layers of width 256 with 4 attention heads, rotary positions and RMSNorm, and differ only in the feed-forward layers.

The dense model has SwiGLU networks of width 768 and 7.2 million parameters in all. The mixture with 8 experts sends every token to 2 experts of width 384: it does the same work per token as the dense model, with 7.2 million active parameters, but holds 17.8 million.

The mixture with 32 experts cuts the same budget finer, 8 of 32 experts of width 96. The fourth model is dense with networks of width 3,072: as many parameters as the mixtures, every one of them used for every token, at about two and a half times the compute.

Each model trained three times, with different seeds, for 3,000 steps of 32 sequences of 256 tokens, about 25 million tokens, on an NVIDIA T4 on Kaggle, with the auxiliary balancing loss described below for the mixtures.

Validation loss during 3,000 training steps, mean over English, Russian and Python, and the final loss on each: Dense 3.076 (English 3.95, Russian 2.99, Python 2.28); Dense, 4x wider 2.958 (English 3.84, Russian 2.83, Python 2.20); MoE, 8 experts 2.989 (English 3.87, Russian 2.88, Python 2.21); MoE, 32 experts 2.988 (English 3.86, Russian 2.90, Python 2.21) Validation loss during 3,000 training steps, mean over English, Russian and Python, and the final loss on each: Dense 3.076 (English 3.95, Russian 2.99, Python 2.28); Dense, 4x wider 2.958 (English 3.84, Russian 2.83, Python 2.20); MoE, 8 experts 2.989 (English 3.87, Russian 2.88, Python 2.21); MoE, 32 experts 2.988 (English 3.86, Russian 2.90, Python 2.21)

At the same compute per token, the mixture with 8 experts ends with a validation loss of 2.985 on average over the three seeds, against 3.062 for the dense model.

The three seeds of the dense model end within 0.024 of each other and those of the mixture within 0.008, so a gap of 0.077 is not noise, and the mixture is better on each kind of text: by 0.08 on English, 0.10 on Russian and 0.06 on Python.

The dense model with networks four times wider, as many parameters as the mixture and all of them working on every token, reaches 2.948, and until about step 1,750 the mixtures are even ahead of it.

The mixture with 8 experts gets two thirds of the way from the small dense model to the large one while computing as much as the small one.

Compute per token is not the same thing as time. Here is how long one training step takes, with the experts run by the loop shown earlier and by the sorted, batched version:

Time of one training step on the Tesla T4: Dense 144 ms; Dense, 4x wider 307 ms; MoE, 8 experts, sorted and batched 188 ms; MoE, 8 experts, loop over experts 189 ms; MoE, 32 experts, sorted and batched 253 ms; MoE, 32 experts, loop over experts 307 ms Time of one training step on the Tesla T4: Dense 144 ms; Dense, 4x wider 307 ms; MoE, 8 experts, sorted and batched 188 ms; MoE, 8 experts, loop over experts 189 ms; MoE, 32 experts, sorted and batched 253 ms; MoE, 32 experts, loop over experts 307 ms

The mixture with 8 experts takes 31% longer than the dense model that does the same arithmetic, and the one with 32 experts 76% longer: routing, sorting and gathering the tokens are work that a count of multiplications doesn’t show.

The dense model with networks four times wider takes 307 milliseconds, more than twice as long as the narrow one. The loop itself turns out to be harmless here.

With 8 experts a batch of 8,192 tokens gives each expert about 2,000 of them, and the loop is as fast as the batched version, 189 against 188 milliseconds per step; with 32 narrower experts it is 21% slower.

Keeping the experts busy

Left alone, the router has a bad habit. Early in training some expert happens to be a little better than the others, the router sends it more tokens, it trains on more data and gets better still, and the router sends it even more.

The other experts get few tokens, learn little and stay unattractive. Shazeer and colleagues, who brought mixtures of experts to large networks, observed this “self-reinforcing” imbalance and countered it with an extra loss[2].

one expert starts a little betterthe router sends it more tokensit trains on more and gets better stillthe other experts learn little

The busiest expert's share of the tokens as a multiple of an even share, mean of 6 layers, during training of the model with 8 experts: no balancing 2.48 at the start and 3.10 at the end; auxiliary loss 1.23 at the start and 1.13 at the end; bias, no loss 1.87 at the start and 1.02 at the end The busiest expert's share of the tokens as a multiple of an even share, mean of 6 layers, during training of the model with 8 experts: no balancing 2.48 at the start and 3.10 at the end; auxiliary loss 1.23 at the start and 1.13 at the end; bias, no loss 1.87 at the start and 1.02 at the end

I trained the mixture with 8 experts in three versions: without any balancing, with the loss described below, and with the bias method that comes after it.

Without balancing, the busiest expert of a layer has 2.6 to 3.1 times its fair share of the tokens by step 200, and nothing ever reverses it.

On held-out text, in three of the six layers of the first of these models a single expert receives 94 to 98% of the tokens, and 9 of the 48 experts get less than a tenth of an even share. They are practically dead: their weights take memory and learn almost nothing.

With the balancing loss, the busiest expert stays at about 1.1 times an even share for the whole run. The unbalanced mixture is still better than the dense model, 3.013 against 3.062, but worse than the balanced one and less predictable: its three seeds end between 2.991 and 3.030.

The balancing loss most models use comes from the Switch Transformer[5]. For a batch of tokens, let \(f_i\) be the share of the token-expert assignments that went to expert \(i\), and \(P_i\) the mean probability the router gave expert \(i\) over the same tokens. The loss is

\[L_{\text{balance}} = \alpha \cdot E \sum_{i=1}^{E} f_i\, P_i ,\]

added to the language modelling loss with a small weight \(\alpha\), 0.01 in the Switch Transformer and in my runs. When the router is perfectly even, \(f_i = P_i = 1/E\) and the sum times \(E\) is 1. When it favours some experts, their large \(f_i\) meet large \(P_i\) and the sum grows.

The two factors play different roles. \(f_i\) is a count of decisions, and a count has no gradient. \(P_i\) is an average of softmax outputs and has one.

So the loss pushes down the probabilities of experts that are already busy, in proportion to how busy they are, and the next tokens drift to the others. Here are both factors in two trained routers, one trained with the loss and one without:

Share of the assignments f and mean router probability P for each expert of layer 3, on held-out text, for two trained routers. no balancing: assignments from 0% to 49% per expert, mean probabilities from 1% to 58%, E times the sum of f·P 2.68; auxiliary loss: assignments from 11% to 16% per expert, mean probabilities from 11% to 15%, E times the sum of f·P 1.00 Share of the assignments f and mean router probability P for each expert of layer 3, on held-out text, for two trained routers. no balancing: assignments from 0% to 49% per expert, mean probabilities from 1% to 58%, E times the sum of f·P 2.68; auxiliary loss: assignments from 11% to 16% per expert, mean probabilities from 11% to 15%, E times the sum of f·P 1.00

In layer 3 of the router trained without the loss, expert 2 gets 49% of the assignments, and the router’s mean probability for it is 0.58, so the router is not hesitating between experts, it is sure of its favourite.

Experts 1 and 8 get almost nothing. The loss would be 2.68 for this router. For the router trained with it, it is 1.00, and its eight experts all get between 11% and 16% of the assignments.

The extra loss has a cost: its gradient flows into the same weights as the language modelling loss and pulls them towards a goal that has nothing to do with predicting text.

Wang and colleagues at DeepSeek proposed balancing without a loss[6], and DeepSeek-V3 uses their method[7]. Every expert gets a bias \(b_i\) that is added to its score when the \(k\) experts are chosen, and only then:

\[\mathcal{T} = \text{the } k \text{ largest of } p_1 + b_1, \dots, p_E + b_E, \qquad \text{weights from } p_i \text{ alone}.\]

After every training step, the bias of each expert that got fewer tokens than the average goes up by a small step \(\gamma\), and the bias of each one that got more goes down by \(\gamma\).

The biases are not learned by gradient descent and don’t touch the model’s weights; they act like a traffic controller that keeps nudging the queues until they are even. I used \(\gamma = 0.001\), the value DeepSeek-V3 used for most of its training.

The biases start at zero and move by 0.001 a step, so they need time: for the first 200 steps the busiest expert gets one and a half to two times its share.

From about step 500 on, the biases keep the load more even than the loss does, 1.02 times an even share at the end against 1.10 to 1.13.

In quality the two are the same at my scale: 2.981 with the biases and 2.985 with the loss, while the seeds of each spread over 0.014 and 0.008.

DeepSeek’s comparison on models with a billion parameters is hardly more decisive: a validation loss of 2.253 with the biases, 2.258 with a balancing loss computed over each sequence, and 2.253 with one computed over the whole batch, as mine is[7].

A fixed budget per expert

Balancing on average is not enough when the experts sit on different GPUs. Each GPU has to reserve memory for the tokens its experts will receive before it knows how many will come, and the shapes of the arrays on an accelerator are fixed in advance.

The Switch Transformer and GShard therefore give every expert a capacity, a fixed number of slots per batch[8]:

\[C = c \cdot \frac{T\, k}{E},\]

where \(T\) is the number of tokens in the batch, \(T k / E\) is what each expert would get with perfectly even routing, and the capacity factor \(c\), usually between 1 and 2, leaves room for unevenness.

Tokens beyond the capacity are dropped by that expert. A dropped token loses that expert’s contribution; if all its experts are full, it passes through the layer unchanged, carried only by the residual connection.

Tokens per expert against the capacity, and the share of assignments dropped for capacity factors from 1 to 2, for three routers of the model with 8 experts. no balancing: the busiest expert gets 3.89 times an even share; dropped 33.3% at 1, 26.3% at 1.25, 21.4% at 1.5, 14.8% at 2; auxiliary loss: the busiest expert gets 1.25 times an even share; dropped 6.1% at 1, 1.0% at 1.25, 0.3% at 1.5, 0.0% at 2; bias, no loss: the busiest expert gets 1.17 times an even share; dropped 5.2% at 1, 0.1% at 1.25, 0.0% at 1.5, 0.0% at 2 Tokens per expert against the capacity, and the share of assignments dropped for capacity factors from 1 to 2, for three routers of the model with 8 experts. no balancing: the busiest expert gets 3.89 times an even share; dropped 33.3% at 1, 26.3% at 1.25, 21.4% at 1.5, 14.8% at 2; auxiliary loss: the busiest expert gets 1.25 times an even share; dropped 6.1% at 1, 1.0% at 1.25, 0.3% at 1.5, 0.0% at 2; bias, no loss: the busiest expert gets 1.17 times an even share; dropped 5.2% at 1, 0.1% at 1.25, 0.0% at 1.5, 0.0% at 2

For the balanced routers a capacity factor of 1 is tight: any unevenness overflows, and 5 to 6% of the assignments are dropped. At 1.25 the share falls to 1% or less, and at 1.5 practically nothing is dropped.

The router trained without balancing is another story. At a capacity factor of 1 it drops a third of the assignments, and even at 2 it drops 15%, because its favourite experts receive nearly four times their share.

The dropped assignments cost more than their share suggests. On the held-out sequences of this analysis, the loss of the mixture with the balancing loss is 2.97 without a limit, 3.21 with 6% of the assignments dropped and 3.01 with 1%. My models were trained without a capacity limit and never learned to do without a missing expert.

A model trained with one adapts: the Switch Transformer saw “typically <1%” of tokens dropped[5], and the ST-MoE study found that models “dropping 10-15% of tokens can perform approximately as well as models that drop <1%”[9].

Dropping tokens is a compromise between quality and the cost of padding, and newer systems avoid it. MegaBlocks computes the experts as one block-sparse product in which every expert gets exactly the rows it was sent[4], and DeepSeek-V3 drops no tokens in training or inference, relying on balancing to keep the load even[7].

What the experts learn

Nobody tells an expert what to specialise in, so it is a fair question whether they specialise at all. My corpus makes it easy to see: a third English, a third Russian, a third Python, three kinds of text that differ in alphabet, vocabulary and structure. If an expert’s tokens are split evenly between the three, it has no preference.

What each expert of the model with 8 experts receives on held-out text with equal amounts of English, Russian and Python. The middle layers are mixed; the first and especially the last layer have experts that lean to one kind of text. Experts that get at least 60% of their tokens from one kind, per layer: layer 1: 1 Russian, 1 Python; layer 2: none; layer 3: none; layer 4: 1 English, 2 Python; layer 5: 2 English, 1 Russian, 2 Python; layer 6: 1 English, 2 Russian, 2 Python What each expert of the model with 8 experts receives on held-out text with equal amounts of English, Russian and Python. The middle layers are mixed; the first and especially the last layer have experts that lean to one kind of text. Experts that get at least 60% of their tokens from one kind, per layer: layer 1: 1 Russian, 1 Python; layer 2: none; layer 3: none; layer 4: 1 English, 2 Python; layer 5: 2 English, 1 Russian, 2 Python; layer 6: 1 English, 2 Russian, 2 Python

The early layers have little preference. In layers 2 and 3 no expert gets more than about half of its tokens from one kind of text, and the preference grows with depth: three experts of layer 4 lean to one kind, five of layer 5.

In layer 6 the experts have divided the languages between them. Five of the eight get 83 to 96% of their tokens from one kind of text, two from Python, two from Russian and one from English; a sixth leans to English with 59%, and two remain mixed.

The first layer stands apart from its neighbours: one of its experts gets 83% of its tokens from Russian text, another 67% from Python.

Here is the same thing token by token, on three short snippets of held-out text. Each token is coloured by the expert it goes to first:

The first-choice expert of every token of three held-out snippets, English, Russian and Python, in the first and the last layer of the trained model with 8 experts. A token is coloured by the kind of text its expert leans to, if at least 60% of that expert's tokens are of one kind. Tokens that go to an expert of their own kind: layer 1: 0 of 13 English, 10 of 17 Russian, 6 of 16 Python; layer 6: 10 of 13 English, 14 of 17 Russian, 9 of 16 Python. In the first layer the Russian expert takes mostly the Cyrillic pieces of names, the English one short words, and a third expert Python's symbols and newlines; in the last, almost every token goes to an expert of its language The first-choice expert of every token of three held-out snippets, English, Russian and Python, in the first and the last layer of the trained model with 8 experts. A token is coloured by the kind of text its expert leans to, if at least 60% of that expert's tokens are of one kind. Tokens that go to an expert of their own kind: layer 1: 0 of 13 English, 10 of 17 Russian, 6 of 16 Python; layer 6: 10 of 13 English, 14 of 17 Russian, 9 of 16 Python. In the first layer the Russian expert takes mostly the Cyrillic pieces of names, the English one short words, and a third expert Python's symbols and newlines; in the last, almost every token goes to an expert of its language

In the first layer the router sorts tokens by what they look like. Its Russian expert takes президент and most of the pieces into which the tokenizer cut the names Либерии, Чарльз and Кинг: Л, иб, Ч, ар, ль, з, Д, Б, инг.

Python’s =, _, (), try and newlines go to an expert that leans to code. English has no expert of its own in this layer at all. At this depth a token’s vector is still mostly the token itself, and the router has little else to go on.

In the last layer, 10 of the 13 English tokens, 14 of the 17 Russian ones and 9 of the 16 Python ones go to an expert of their own kind.

All the others but one go to the two mixed experts of this layer, and most of them are tokens the three kinds of text share: the numbers 3 and 21, the full stops after the initials Д. and Б., and in Python the brackets, colons, the underscore and newlines.

Key point

In the last layer the router sorts tokens by language, because that layer prepares the prediction of the next token, and what the next token can be depends first of all on the language.

Specialisation by topic, the kind people usually hope for, is much weaker in large models. Mistral looked for experts on topics in Mixtral and did not find them: “the distribution of expert assignment is very similar for ArXiv papers (written in Latex), for biology (PubMed Abstracts), and for Philosophy (PhilPapers) documents”.

What they found instead was syntax: “in code, the indentation tokens are always assigned to the same experts, particularly at the first and last layers where the hidden states are more correlated to the input and output of the model”[10].

In OLMoE, trained from scratch with 64 small experts per layer, the specialisation is stronger: “for arXiv, which has a very specific distribution with lots of scientific text, the first expert in layer 0 is nearly 100% specialized”, and “vocabulary specialization is higher in later layers”[11].

Languages are a case of their own. The ST-MoE study found that “experts trained on multilingual mixtures do not exhibit language specialization” and blamed the balancing: a group of tokens over which the loss was computed held only one or a few languages, so every expert had to take every language to keep the load even within the group[9].

My batches mix all three kinds of text, and an English expert can stay busy on the English part of every batch.

Smaller experts, and one for everybody

An expert in Mixtral[10] is as wide as the whole feed-forward network of Mistral 7B, whose architecture Mixtral otherwise shares, and a token picks 2 of 8.

DeepSeekMoE argued that this is too coarse: with only 8 experts, each has to cover a broad mix of tokens, and knowledge that all of them need is learned several times over[12]. Their remedy has two parts.

The first is to cut every expert into \(m\) smaller ones and route each token to \(m\) times as many, which keeps the parameters and the compute the same and multiplies the number of combinations a token can choose from: in their example, top-2 of 16 experts allows 120 combinations, top-8 of 64 quarter-size experts “4,426,165,368 potential combinations”.

The second is to set aside shared experts that every token passes through, for what all tokens need, so that the routed experts can specialise.

Three ways to cut the same budget of experts. Eight experts with two per token can be combined in 28 ways; 32 experts a quarter of the size with eight per token, the same compute and the same parameters, in 10,518,300 ways. DeepSeek-V3 has one shared expert every token passes through and 256 routed experts with eight per token: about 4.1e+14 combinations Three ways to cut the same budget of experts. Eight experts with two per token can be combined in 28 ways; 32 experts a quarter of the size with eight per token, the same compute and the same parameters, in 10,518,300 ways. DeepSeek-V3 has one shared expert every token passes through and 256 routed experts with eight per token: about 4.1e+14 combinations

In my models the finer cut made no difference. The mixture with 32 experts ends at 2.983 and the one with 8 at 2.985, well within the spread of the seeds, and the finer one costs more time, 253 milliseconds per step against 188, because routing, sorting and gathering grow with the number of experts. It does take dropped tokens better.

With a capacity factor of 1 it loses the same 6% of its assignments as the coarse mixture, but its loss rises by 0.10 instead of 0.24: a dropped assignment takes away one of a token’s eight experts instead of one of two.

The case for finer experts was made on larger models. DeepSeekMoE’s smallest comparison is at 2 billion parameters, where “DeepSeekMoE 2B achieves comparable performance with GShard 2.9B”[12], a hundred times the size of my models, trained on far more than my 25 million tokens.

Mixtures in real models

Here are the experts of open models as their config files describe them, in the order the models came out:

Routed experts per token and per layer, and the active share of parameters, in open mixture-of-experts models in the order of release: Mixtral 8x7B (Dec 2023): 2 of 8 experts, 12.9B of 46.7B parameters active; DeepSeek-V2 (May 2024): 6 of 160 experts plus 2 shared, 21B of 236B parameters active; OLMoE-1B-7B (Sep 2024): 8 of 64 experts, 1.3B of 6.9B parameters active; DeepSeek-V3 (Dec 2024): 8 of 256 experts plus 1 shared, 37B of 671B parameters active; Llama 4 Maverick (Apr 2025): 1 of 128 experts plus 1 shared, 17B of 400B parameters active; Qwen3-235B-A22B (Apr 2025): 8 of 128 experts, 22B of 235B parameters active; Kimi K2 (Jul 2025): 8 of 384 experts plus 1 shared, 32B of 1,000B parameters active; GLM-4.5 (Jul 2025): 8 of 160 experts plus 1 shared, 32B of 355B parameters active; gpt-oss-120b (Aug 2025): 4 of 128 experts, 5.1B of 116.8B parameters active; Qwen3-Next-80B-A3B (Sep 2025): 10 of 512 experts plus 1 shared, 3B of 80B parameters active; MiniMax-M2 (Oct 2025): 8 of 256 experts, 10B of 230B parameters active; Qwen3.5-397B-A17B (Feb 2026): 10 of 512 experts plus 1 shared, 17B of 397B parameters active; DeepSeek-V4-Pro (Apr 2026): 6 of 384 experts plus 1 shared, 49B of 1,600B parameters active; Kimi K3 (Jun 2026): 16 of 896 experts plus 2 shared, 104B of 2,800B parameters active Routed experts per token and per layer, and the active share of parameters, in open mixture-of-experts models in the order of release: Mixtral 8x7B (Dec 2023): 2 of 8 experts, 12.9B of 46.7B parameters active; DeepSeek-V2 (May 2024): 6 of 160 experts plus 2 shared, 21B of 236B parameters active; OLMoE-1B-7B (Sep 2024): 8 of 64 experts, 1.3B of 6.9B parameters active; DeepSeek-V3 (Dec 2024): 8 of 256 experts plus 1 shared, 37B of 671B parameters active; Llama 4 Maverick (Apr 2025): 1 of 128 experts plus 1 shared, 17B of 400B parameters active; Qwen3-235B-A22B (Apr 2025): 8 of 128 experts, 22B of 235B parameters active; Kimi K2 (Jul 2025): 8 of 384 experts plus 1 shared, 32B of 1,000B parameters active; GLM-4.5 (Jul 2025): 8 of 160 experts plus 1 shared, 32B of 355B parameters active; gpt-oss-120b (Aug 2025): 4 of 128 experts, 5.1B of 116.8B parameters active; Qwen3-Next-80B-A3B (Sep 2025): 10 of 512 experts plus 1 shared, 3B of 80B parameters active; MiniMax-M2 (Oct 2025): 8 of 256 experts, 10B of 230B parameters active; Qwen3.5-397B-A17B (Feb 2026): 10 of 512 experts plus 1 shared, 17B of 397B parameters active; DeepSeek-V4-Pro (Apr 2026): 6 of 384 experts plus 1 shared, 49B of 1,600B parameters active; Kimi K3 (Jun 2026): 16 of 896 experts plus 2 shared, 104B of 2,800B parameters active

Three things changed over these two and a half years. The experts multiplied and shrank. Mixtral has 8 experts of width 14,336 per layer, DeepSeek-V3 256 of width 2,048, Kimi K3 896.

Moonshot measured the effect directly for Kimi K2: at a fixed number of active parameters, “increasing the total number of experts (i.e., increasing sparsity) consistently lowers both the training and validation loss”, and they settled on 8 of 384 as a balance between quality and the cost of the infrastructure[13].

Krajewski and colleagues reached the same conclusion from scaling laws: setting the size of experts to match the dense feed-forward layer “is not optimal at almost any computational budget”[14].

The active share fell accordingly, from 28% in Mixtral to 3–5% in the models of the past year. A model now holds twenty to thirty times as many parameters as it uses for a token.

The shared expert did not convince everyone. DeepSeek, Moonshot, Zhipu and the newer Qwen models use one or two; the OLMoE team found that “sharing an expert performs slightly worse” and left it out[11], and Qwen3 also has none[15].

Not every layer is a mixture either. DeepSeek-V3 keeps its first three layers dense, and Llama 4 Maverick alternates dense and mixture layers, with a token going to 1 of 128 routed experts plus a shared one[16].

Serving a model that is mostly idle

The first cost of a mixture is memory. Every expert has to be at hand, because any token may need any of them, so DeepSeek-V3 needs room for all 671 billion parameters, about 671 GB at one byte per parameter in FP8, more than the 640 GB of a server with eight 80 GB GPUs.

OpenAI made gpt-oss-120b fit on one GPU by storing the experts, which hold “90+% of the total parameter count”, in a 4-bit format, MXFP4: “quantizing these to MXFP4 enables the larger model to fit on a single 80GB GPU”[17].

A model with all its experts in memory reads only a few of them for each token. At a batch of one, the case of a chatbot on a laptop, a decoding step reads the attention weights and \(k\) experts per layer and nothing else.

Decoding is limited by how fast weights come from memory, not by arithmetic, so the step costs about as much as in a dense model of the active size.

This is why models like Qwen3-30B-A3B are popular on laptops: they need the memory of a 30-billion-parameter model but read only about 3 billion parameters per token. llama.cpp even has options to keep the experts in the computer’s main memory while the rest of the model sits on the GPU, --cpu-moe for all of them and --n-cpu-moe N for the first N layers[18]; the GPU then holds what every token needs, and each token reads only its few experts from the slower memory.

A server decodes many conversations at once, and every conversation’s token picks its own experts. If the router spreads tokens evenly, a given expert is skipped by one token with probability \(1 - k/E\), and by all \(B\) tokens of a batch with probability \((1 - k/E)^B\). The expected share of experts a decoding step has to read is therefore

\[1 - \left(1 - \frac{k}{E}\right)^{B}.\]

Expected share of the experts of a layer that one decoding step has to read, for batches of 1 to 256 sequences, if the router spreads tokens evenly: Mixtral 8x7B (2 of 8) 25% at 1, 99% at 16, 100% at 256; Qwen3-235B-A22B (8 of 128) 6% at 1, 64% at 16, 100% at 256; DeepSeek-V3 (8 of 256) 3% at 1, 40% at 16, 100% at 256; Kimi K2 (8 of 384) 2% at 1, 29% at 16, 100% at 256. Our trained router with 8 of 32 experts, measured: 25% at 1, 68% at 4, 99% at 16, 100% at 64, 100% at 256 Expected share of the experts of a layer that one decoding step has to read, for batches of 1 to 256 sequences, if the router spreads tokens evenly: Mixtral 8x7B (2 of 8) 25% at 1, 99% at 16, 100% at 256; Qwen3-235B-A22B (8 of 128) 6% at 1, 64% at 16, 100% at 256; DeepSeek-V3 (8 of 256) 3% at 1, 40% at 16, 100% at 256; Kimi K2 (8 of 384) 2% at 1, 29% at 16, 100% at 256. Our trained router with 8 of 32 experts, measured: 25% at 1, 68% at 4, 99% at 16, 100% at 64, 100% at 256

My router with 8 of 32 experts sits exactly on its curve: a decoding step needs 25% of the experts for one sequence, 90% for 8 and all of them from 32 on, as if every token picked its experts at random.

A model with a smaller share per token holds out longer. DeepSeek-V3, with 8 of 256, reads 3% of its routed experts for one sequence, 40% for 16 and 87% for 64, and from about 128 sequences on, nearly all of them.

Key point

The saving in memory traffic that makes a mixture fast at a batch of one disappears at the batch sizes a server runs at. What remains is the saving in arithmetic: each expert multiplies only the tokens that chose it.

To turn this saving into speed, a server needs enough tokens per expert to keep each product efficient, and that is what expert parallelism is for: the experts are spread over many GPUs, every GPU holds a few of them, and tokens travel to their experts and back.

Expert parallelism on four GPUs, each holding two of the eight experts and a quarter of the tokens. The tokens of GPU 0 are followed: after the router, each is sent to the GPUs that hold its two experts, in an all-to-all exchange in which every GPU sends to every other; the experts process what they received; a second all-to-all brings the results back to GPU 0, where each token's two results are added Expert parallelism on four GPUs, each holding two of the eight experts and a quarter of the tokens. The tokens of GPU 0 are followed: after the router, each is sent to the GPUs that hold its two experts, in an all-to-all exchange in which every GPU sends to every other; the experts process what they received; a second all-to-all brings the results back to GPU 0, where each token's two results are added

DeepSeek described its production setup for V3 and R1 in early 2025. Its reasoning is the same: “the model’s high sparsity necessitates an extremely large overall batch size. This ensures sufficient batch size per expert, enabling higher throughput and lower latency. Large-scale cross-node EP is essential.”

Reading a prompt, one deployment unit spans 4 nodes with 32-way expert parallelism; decoding, it spans 18 nodes with 144-way expert parallelism, each GPU holding 2 routed experts and the shared one, and each node of 8 GPUs writes about 14,800 tokens per second[19].

The price is communication. Every mixture layer sends every token to other GPUs and back, twice per layer and 58 layers per token in DeepSeek-V3, and the exchange has to be hidden behind computation.

DeepSeek overlaps the communication of one batch with the computation of another and, for training, limits every token to experts on at most 4 nodes, because NVLink inside a node is about 3.2 times faster than InfiniBand between nodes[7].

When a mixture is the right choice

A mixture of experts trades memory for compute. For a given quality it needs fewer operations per token than a dense model, and in exchange it keeps all its experts in memory and, served at scale, moves every token between GPUs and back in every layer.

The gain in compute is real and large. The Switch Transformer reached the quality of a dense T5-Base “in one-seventh the time” at the same compute per token[5].

The gain in wall-clock time is smaller, because routing and moving tokens cost time that no operation count shows: the OLMoE team found that their mixture “reaches the performance of the dense model with ∼3× fewer tokens”, but in time “only ∼2× faster”[11].

So the trade pays where compute is the bottleneck and memory is plentiful: training large models, and serving them to many users at once on clusters where every expert sees enough tokens to stay busy. It pays less where memory is the bottleneck or traffic is light.

The authors of GLaM, one of the first large mixtures, noted that sparse models “require a larger number of devices”, which “increases the serving cost especially when the serving traffic is low”[20].

In practice

Fine-tuning needs care too. In the ST-MoE study, sparse models were “prone to overfitting”: on a task with 250 training examples the sparse model did worse than the dense one, and sparse models needed their own fine-tuning settings, smaller batches and a higher learning rate[9].

A laptop with plenty of memory and a modest processor is a case of its own. Its memory can hold all the experts, and at a batch of one each token reads only the few it needs.

References

  1. Jacobs, R. A., Jordan, M. I., Nowlan, S. J., & Hinton, G. E. (1991). Adaptive mixtures of local experts. Neural Computation, 3(1), 79–87. ↩
  2. Shazeer, N., Mirhoseini, A., Maziarz, K., Davis, A., Le, Q., Hinton, G., & Dean, J. (2017). Outrageously large neural networks: the sparsely-gated mixture-of-experts layer. ICLR 2017. ↩
  3. Geva, M., Schuster, R., Berant, J., & Levy, O. (2021). Transformer feed-forward layers are key-value memories. EMNLP 2021. ↩
  4. Gale, T., Narayanan, D., Young, C., & Zaharia, M. (2023). MegaBlocks: efficient sparse training with mixture-of-experts. MLSys 2023. ↩
  5. Fedus, W., Zoph, B., & Shazeer, N. (2022). Switch Transformers: scaling to trillion parameter models with simple and efficient sparsity. Journal of Machine Learning Research, 23(120), 1–39. ↩
  6. Wang, L., Gao, H., Zhao, C., Sun, X., & Dai, D. (2024). Auxiliary-loss-free load balancing strategy for mixture-of-experts. arXiv. ↩
  7. DeepSeek-AI (2024). DeepSeek-V3 technical report. arXiv. ↩
  8. Lepikhin, D., Lee, H., Xu, Y., Chen, D., Firat, O., Huang, Y., Krikun, M., Shazeer, N., & Chen, Z. (2021). GShard: scaling giant models with conditional computation and automatic sharding. ICLR 2021. ↩
  9. Zoph, B., Bello, I., Kumar, S., Du, N., Huang, Y., Dean, J., Shazeer, N., & Fedus, W. (2022). ST-MoE: designing stable and transferable sparse expert models. arXiv. ↩
  10. Jiang, A. Q., et al. (2024). Mixtral of experts. arXiv. ↩
  11. Muennighoff, N., et al. (2025). OLMoE: open mixture-of-experts language models. ICLR 2025. ↩
  12. Dai, D., et al. (2024). DeepSeekMoE: towards ultimate expert specialization in mixture-of-experts language models. ACL 2024. ↩
  13. Kimi Team (2025). Kimi K2: open agentic intelligence. arXiv. ↩
  14. Krajewski, J., et al. (2024). Scaling laws for fine-grained mixture of experts. ICML 2024. ↩
  15. Qwen Team (2025). Qwen3 technical report. arXiv. ↩
  16. Meta (2025). The Llama 4 herd: the beginning of a new era of natively multimodal AI innovation. Blog post. ↩
  17. OpenAI (2025). gpt-oss-120b & gpt-oss-20b model card. arXiv. ↩
  18. llama.cpp (2025). llama : add –n-cpu-moe option. Pull request #15077. ↩
  19. DeepSeek-AI (2025). Day 6: one more thing, DeepSeek-V3/R1 inference system overview. GitHub. ↩
  20. Du, N., et al. (2022). GLaM: efficient scaling of language models with mixture-of-experts. ICML 2022. ↩
]]> Eva Koroleva Hybrid attention: why new LLMs keep full attention in only a few layers 2026-01-15T12:00:00+00:00 2026-11-02T00:00:00+00:00 https://blog.korolevalab.com/blog/2026/hybrid-attention

In short

In Qwen3-Next, Kimi Linear and the other hybrids, most layers are recurrent: they read the text into a state of fixed size, keep no cache and write every token at the same cost. The price is exact recall, because a \(d \times d\) state holds at most \(d\) key-value pairs exactly.

A few layers of full attention bring the lookups back, though where they sit matters: in a small recall test the same four layers recalled every value with attention second and far fewer with it last.

What a hybrid saves is the cache, four to ten times over, and that pays off where contexts are long and many conversations share a GPU.

3 GiBcache of Qwen3-Next for one 128K-token conversation, against 32 GiB for Qwen3-32B
36%of the values recalled at 32 pairs by DeltaNet, the best purely recurrent model
16×faster to write a token after 65K tokens, one recurrent layer against attention

Qwen3-32B, a dense model from April 2025, keeps 32 GiB of keys and values for one conversation of 128 thousand tokens, about half of what its weights take. Qwen3-Next, released by the same team a few months later, keeps 3 GiB at the same length. It doesn’t compress its cache or skip tokens.

Three of every four of its layers simply have no cache: they read the text the way a recurrent network does, into a state of fixed size, and only every fourth layer is ordinary attention.

Qwen3-Next is not the only one. Jamba, the Nemotron-H models, Kimi Linear and, for a while, MiniMax made the same choice: full attention in a few layers, something cheaper in the rest.

This post builds the cheaper part from scratch, from linear attention to the gated delta rule, tests what it can and cannot remember, and then looks at how real models combine the two and what that buys on a GPU.

What attention remembers

Every attention layer of a Transformer answers the same question for every token: which of the earlier tokens matter here, and what do they say. Softmax attention answers it by brute force. Each token leaves behind a key, which says what it is about, and a value, which says what it contributes.

A new token arrives with a query, compares it with every key, turns the scores into weights with a softmax, and takes the weighted sum of the values:

\[o_t = \sum_{j \le t} \frac{\exp(q_t \cdot k_j / \sqrt{d})}{\sum_{i \le t} \exp(q_t \cdot k_i / \sqrt{d})}\, v_j .\]

Nothing is ever thrown away. To write token \(t\), the layer needs the keys and values of all \(t-1\) tokens before it, which is why serving engines keep them in a KV cache instead of recomputing them.

Memory therefore grows linearly with the text, and since each new token reads all of it, reading a whole text of \(n\) tokens takes time proportional to \(n^2\).

Think of a library that keeps every book it has ever received and checks the whole catalogue for every question. The answers are exact, but the building has to grow with the collection.

Attention without the softmax

The softmax is what forces the rereading: the exponential ties every query to every key. Katharopoulos and colleagues replaced it, putting a dot product of features \(\phi(q) \cdot \phi(k)\) in place of \(\exp(q \cdot k)\), with a positive map such as \(\phi(x) = \mathrm{elu}(x) + 1\), and noticed that the rest follows from associativity[1]:

\[o_t = \frac{\sum_{j \le t} \big(\phi(q_t)^\top \phi(k_j)\big) v_j}{\sum_{j \le t} \phi(q_t)^\top \phi(k_j)} = \frac{\phi(q_t)^\top S_t}{\phi(q_t)^\top z_t}, \qquad S_t = \sum_{j \le t} \phi(k_j)\, v_j^\top, \quad z_t = \sum_{j \le t} \phi(k_j).\]

The query no longer sits inside the sum. Everything the layer needs from the past now sits in \(S_t\), a \(d \times d\) matrix, and \(z_t\), a vector, and both are updated by adding one term per token:

\[S_t = S_{t-1} + \phi(k_t)\, v_t^\top .\]

This is a recurrent network. Its state has the same size after ten tokens and after ten million, writing a token costs the same at any position, and there is no cache. Done for all positions at once, the same computation is the attention formula with the brackets moved:

The same product computed in two orders, for 32,768 tokens and 128-dimensional heads. Multiplying the queries by the keys first gives a 32,768 by 32,768 matrix of scores, 1,073,741,824 numbers, before the values are applied. Multiplying the keys by the values first gives a 128 by 128 matrix, 16,384 numbers, that every query then reads. Without the softmax between them, the order is free The same product computed in two orders, for 32,768 tokens and 128-dimensional heads. Multiplying the queries by the keys first gives a 32,768 by 32,768 matrix of scores, 1,073,741,824 numbers, before the values are applied. Multiplying the keys by the values first gives a 128 by 128 matrix, 16,384 numbers, that every query then reads. Without the softmax between them, the order is free

Both forms compute the same numbers, only in a different order. In code, the recurrent form loops over positions and the parallel one builds the masked score matrix:

def linear_recurrent(q, k, v):
    q, k = feature(q), feature(k)
    state = q.new_zeros(*q.shape[:2], q.shape[-1], v.shape[-1])  # S = sum of k v^T
    norm = q.new_zeros(*q.shape[:2], q.shape[-1])                # z = sum of k
    out = []
    for t in range(q.shape[2]):
        state = state + k[:, :, t, :, None] * v[:, :, t, None, :]
        norm = norm + k[:, :, t]
        out.append(torch.einsum("bhk,bhkv->bhv", q[:, :, t], state)
                   / torch.einsum("bhk,bhk->bh", q[:, :, t], norm)[..., None])
    return torch.stack(out, dim=2)

def linear_attention(q, k, v):
    q, k = feature(q), feature(k)
    causal = torch.ones(q.shape[2], q.shape[2], dtype=torch.bool).tril()
    scores = (q @ k.transpose(-1, -2)).masked_fill(~causal, 0)
    return (scores @ v) / scores.sum(-1, keepdim=True)

On random inputs in double precision the two functions agree to about \(10^{-15}\), which is rounding error. The recurrent form is how the layer writes text, one token at a time; the parallel form is how it trains, all positions at once.

The parallel form also shows a catch: it still builds an \(n \times n\) matrix of scores, because with a causal mask every position needs its own prefix of the sum and the brackets can’t simply move. Training a linear layer efficiently needs a third form, which comes later.

A memory of fixed size

The price of the fixed state can be read straight from the formula. The state is one matrix. Every pair \((k_j, v_j)\) is added into the same \(d^2\) numbers, and to recall the value stored under a key, the layer multiplies the state by it:

\[S^\top k_i = \sum_j (k_j \cdot k_i)\, v_j = (k_i \cdot k_i)\, v_i + \sum_{j \ne i} (k_j \cdot k_i)\, v_j .\]

The first term is the stored value. The second is interference: every other pair, weighted by how similar its key is to \(k_i\). If the keys are orthogonal, the interference is exactly zero, but a \(d\)-dimensional space has room for only \(d\) orthogonal directions.

Beyond that, each new pair blurs the ones already stored. The library has turned into one notebook page on which every note is written over the others; it stays readable while there are few notes.

Softmax attention has no such limit because it never mixes the notes: it keeps them all and lets the exponential pick out the right one. It pays for that with the growing cache. Everything below is about this trade between a fixed state and exact recall.

Arora and colleagues measured it across efficient architectures and called it “a key tradeoff between a model’s state size and recall ability”[2]: the smaller the state a model keeps while reading, the worse it recalls.

Writing with the delta rule

Plain addition has a second, quieter problem. If the same key comes back with a new value, say a variable that has been reassigned, the state holds the sum of the old and the new value, and reading it returns a mixture of both. Nothing in the notebook can be crossed out.

The remedy dates from 1960, when Widrow and Hoff trained their linear neurons with what became known as the delta rule: compare the output with the target and move the weights in proportion to the error[3].

Schlag, Irie and Schmidhuber noticed that the state of linear attention is a weight matrix, written at every step, and proposed writing it the same way[4]. Before storing \(v_t\) under \(k_t\), read what is already there, and write only the difference:

\[S_t = S_{t-1} + \beta_t\, k_t \big(v_t - S_{t-1}^\top k_t\big)^\top = (I - \beta_t k_t k_t^\top)\, S_{t-1} + \beta_t\, k_t v_t^\top ,\]

with \(k_t\) of unit length and \(\beta_t \in (0, 1)\), a writing strength the layer computes for each token. This is one step of gradient descent on the memory’s error for this key, \(\tfrac12 \lVert S^\top k_t - v_t \rVert^2\), with learning rate \(\beta_t\): the gradient with respect to \(S\) is \(k_t (S^\top k_t - v_t)^\top\).

With \(\beta_t = 1\), the old value under \(k_t\) is erased and the new one is stored exactly. A layer that writes this way is called DeltaNet. In the space of values, one write looks like this:

One write of the delta rule, drawn in the space of values. The memory returns an old value for the key; the new value is elsewhere; the error is the arrow between them. A write of strength 0.5 moves the memory's answer halfway along the error, a write of strength 1 lands on the new value exactly. Plain addition would instead add the new value to the old one One write of the delta rule, drawn in the space of values. The memory returns an old value for the key; the new value is elsewhere; the error is the arrow between them. A write of strength 0.5 moves the memory's answer halfway along the error, a write of strength 1 lands on the new value exactly. Plain addition would instead add the new value to the old one

Side by side, the three memories differ in what they keep and in how a new pair changes it:

Three ways for a layer to remember the tokens it has read. Softmax attention keeps the key and the value of every token, so its memory grows with the text, and a query is compared with every key. Linear attention keeps one matrix and adds the outer product of each key and value to it, so its memory has a fixed size. The delta rule keeps the same matrix but first reads what it holds under the new key and writes only the difference Three ways for a layer to remember the tokens it has read. Softmax attention keeps the key and the value of every token, so its memory grows with the text, and a query is compared with every key. Linear attention keeps one matrix and adds the outer product of each key and value to it, so its memory has a fixed size. The delta rule keeps the same matrix but first reads what it holds under the new key and writes only the difference

How much of a difference does editing make? A small experiment shows both sides of it: a \(64 \times 64\) memory, random unit keys and values, written once by adding and once by the delta rule, and read back by finding the nearest stored value. Softmax attention, which keeps every pair, is the reference.

Left: share of 64-dimensional key-value pairs recalled from a 64 by 64 memory. Keeping every pair, as softmax attention does, recalls all of them at any number. Adding outer products recalls 98% of 128 pairs and 34% of 512; editing with the delta rule recalls 53% and 13%, keeping the newest. Right: after 32 pairs, some keys get a new value; the delta rule returns the new value 100% of the time, adding and softmax attention about half the time Left: share of 64-dimensional key-value pairs recalled from a 64 by 64 memory. Keeping every pair, as softmax attention does, recalls all of them at any number. Adding outer products recalls 98% of 128 pairs and 34% of 512; editing with the delta rule recalls 53% and 13%, keeping the newest. Right: after 32 pairs, some keys get a new value; the delta rule returns the new value 100% of the time, adding and softmax attention about half the time

The left panel holds a surprise: from 64 pairs on, plain addition recalls more than the delta rule. With \(\beta_t = 1\), each write makes the memory return exactly \(v_t\) for \(k_t\), and since random keys are not orthogonal, it also shifts what the memory returns for every older key that overlaps with \(k_t\).

Once the memory is full, the delta rule holds the recent pairs well and the old ones poorly. Addition keeps a blur of everything, in which the nearest value is still often the right one.

The right panel shows what the delta rule is for. After 32 pairs, some of the keys get a new value. The delta rule returns the new value nearly every time.

Addition returns the sum of the old and the new value, which is nearer to the new one only about half the time, and softmax attention as written here does no better: it sees two identical keys and averages their values.

A trained Transformer can tell the later pair from the earlier one by their positions; this toy has no positions.

Key point

The delta rule doesn’t enlarge the memory. A \(d \times d\) matrix stays a \(d \times d\) matrix, and it holds at most \(d\) pairs exactly, whichever way they are written. What the delta rule adds is the ability to correct what is stored, which plain addition doesn’t have at all.

Forgetting on purpose

Editing handles values that change. It does nothing about the opposite problem: a long text moves on, and much of what the state holds stops mattering. A memory of fixed size needs a way to let go of it.

That is what gates are for. RetNet multiplied the state by a fixed decay at every step[5], Gated Linear Attention made the decay depend on the token[6], and Mamba-2 showed that its own selective state space model is linear attention with a data-dependent scalar gate[7]. Gated DeltaNet puts the gate and the delta rule together[8]:

\[S_t = \alpha_t \,(I - \beta_t k_t k_t^\top)\, S_{t-1} + \beta_t\, k_t v_t^\top, \qquad \alpha_t \in (0, 1).\]

Both \(\alpha_t\) and \(\beta_t\) are computed from the current token, so the layer decides as it reads when to forget and how firmly to write.

The two parts split the work: in the authors’ words, the gate “can promptly clear memory by setting \(\alpha_t \to 0\)”, and with \(\alpha_t \to 1\) the layer is “effectively switching to the pure delta rule” and edits single entries without touching the rest.

This is the layer Qwen3-Next is built on, and Kimi Delta Attention is the same update with a separate gate for every channel of the state instead of one per head[9].

Training a recurrence in parallel

A recurrence over positions is awkward to train on a GPU, which is fast at large matrix products and slow at loops of ten thousand small updates. The solution is to cut the sequence into chunks.

Inside a chunk of 64 tokens, the writes depend on each other only through a small lower-triangular system that can be solved with matrix products; between chunks, only the \(d \times d\) state is passed along[10].

How a recurrent layer is trained in parallel. The sequence is cut into chunks of 64 tokens. Inside a chunk, all positions are computed at once with matrix products; between chunks, only the d by d state is passed on, one step per chunk instead of one per token How a recurrent layer is trained in parallel. The sequence is cut into chunks of 64 tokens. Inside a chunk, all positions are computed at once with matrix products; between chunks, only the d by d state is passed on, one step per chunk instead of one per token

The chunked form is exact. My plain-PyTorch version of it agrees with the token-by-token loop to \(1.4 \times 10^{-15}\) in double precision, with the gate or without it, and it is what the experiments below train with. Production models use the Triton kernels of flash-linear-attention, which do the same thing much faster.

A test of memory: associative recall

The language modelling loss says little about memory, because much of predicting the next word needs little of it. A more direct test asks a model to remember what it was told. In multi-query associative recall, a sequence lists key-value pairs and then asks for the values of the keys in a shuffled order[11]:

k1 v1  k2 v2  ...  kN vN   |   q1 a1  q2 a2  ...  qN aN

Each \(q_i\) is one of the keys, and the model must put the right value after it. It is the skill behind answering a question about a long document, using a name defined forty pages earlier, or calling a function with the arguments described at the top of a prompt.

I trained six small models on it, each with four layers of width 64 and two heads, the same in everything except how their layers mix tokens: softmax attention in all four, linear attention, DeltaNet, Gated DeltaNet, and two hybrids with three Gated DeltaNet layers and one softmax attention layer.

In one hybrid the attention layer is second; in the other it is on top, the place it has in each block of four in Qwen3-Next.

Every layer, attention included, first passes its queries, keys and values through a short convolution over the last four positions, as Mamba and DeltaNet do, so that a value can carry its key with it.

Each model saw 64 fresh random sequences at each of 4,000 steps and was scored on held-out ones. The models with Gated DeltaNet layers trained on an Apple M2, the others on a T4 (code).

Share of values recalled with 4, 8, 16, 32 and 64 key-value pairs, by six models of four layers and width 64: Transformer 100%, 100%, 100%, 100%, 100%; Hybrid (attention 2nd) 100%, 100%, 100%, 100%, 100%; Hybrid (attention on top) 100%, 100%, 26%, 12%, 1%; Gated DeltaNet 100%, 98%, 19%, 2%, 0%; DeltaNet 100%, 100%, 71%, 36%, 13%; Linear attention 97%, 82%, 76%, 26%, 0% Share of values recalled with 4, 8, 16, 32 and 64 key-value pairs, by six models of four layers and width 64: Transformer 100%, 100%, 100%, 100%, 100%; Hybrid (attention 2nd) 100%, 100%, 100%, 100%, 100%; Hybrid (attention on top) 100%, 100%, 26%, 12%, 1%; Gated DeltaNet 100%, 98%, 19%, 2%, 0%; DeltaNet 100%, 100%, 71%, 36%, 13%; Linear attention 97%, 82%, 76%, 26%, 0%

The Transformer recalls every value at every size on the chart. The three recurrent models hold up while there are few pairs and fall apart between 8 and 32: at 32 pairs DeltaNet recalls 36% of the values, linear attention 26% and Gated DeltaNet 2%, and at 64 the best of them, DeltaNet, recalls 13%.

Each head here works with keys of 32 numbers, so by the bound from earlier its state holds at most 32 pairs exactly, and at 16 pairs all three are already well short of full recall.

At 128 pairs even the Transformer didn’t learn the task in 4,000 steps, which is why the chart stops at 64.

Gated DeltaNet, the newest of the three, comes last from 16 pairs on. That is less strange than it looks, because this test gives the gate nothing to do. Every pair has to be kept until its question comes, so forgetting can only lose information.

The authors of Gated DeltaNet, who trained models of 1.3 billion parameters on 100 billion tokens of real text, found that it “consistently surpasses existing models like Mamba2 and DeltaNet”, in-context retrieval included[8]. Four layers and 4,000 steps of a synthetic task are a different regime.

Trap

The gate is the easiest part to get wrong. In my first version it started at about one half, and weight decay kept pulling it back there, so the state lost half of what it held at every token, and Gated DeltaNet and the hybrid stayed at chance level. Starting the gate as the Gated DeltaNet and Mamba-2 code does, where a typical head keeps about 94% of its state from one token to the next, and leaving it out of weight decay, fixed it.

The two hybrids differ only in the order of their layers. With attention second, the hybrid recalls every value at every size, like the Transformer, though three of its four layers are Gated DeltaNet, which on its own recalls the least.

With attention on top, the same layers recall 26% of the values at 16 pairs and 12% at 32, at the better of two learning rates: a little more than Gated DeltaNet alone and less than DeltaNet.

I can’t say for certain why. The attention layer on top has what it needs: the convolution hands each value its key, and the other hybrid shows that one attention layer among three Gated DeltaNet layers is enough. This one just doesn’t learn the lookup in 4,000 steps.

A model of four layers is far from Qwen3-Next, where twelve attention layers are spread through 48 and most of them have many recurrent layers on both sides, so this says nothing about whether its layout is a good one.

The ratio alone doesn’t describe a hybrid: with the same layers in a different order, one model recalled every value, like a Transformer, and the other 12% of them at 32 pairs, like a recurrent network.

Why one layer of attention is enough

Attention solves associative recall in one layer: the query for \(k_i\) meets the stored key \(k_i\), the softmax puts nearly all the weight on it, and out comes \(v_i\), however many pairs came before.

Arora and colleagues proved that attention needs a model dimension independent of the sequence length for this, while the gated convolutions they studied need one that grows with it.

Measuring the gap in quality between the two kinds of models on real text, they found that “82% of the gap is explained by each model’s ability to recall information that is previously mentioned in-context”[11].

Copying gives the same picture. Jelassi and colleagues proved that two layers of attention can copy strings of exponential length, while models with a fixed state are “fundamentally limited by their fixed-size latent state”[12].

None of this requires attention everywhere. Much of predicting the next token is local, or a matter of keeping a running summary of what the text is about, and a gated recurrent layer does that well and cheaply.

Exact lookups into the distant past are rarer, and a layer that can make them serves the whole model: the recurrent layers around it carry the local context, the attention layer finds the match. In the library picture, most floors work from a notebook page, and every few floors someone is allowed into the archive.

The authors of Mamba-2 tried such hybrids and described the division of labour the same way:

Attention layers act as a retrieval mechanism to quickly refer to previous tokens in the sequence instead of forcing the model to compress all the context to its memory.

Tri Dao and Albert Gu, the authors of Mamba-2

They found that “around 10% of the total number of layers being attention performs best”[7].

A study at NVIDIA trained an 8-billion-parameter Transformer, pure Mamba and Mamba-2 models and a hybrid with 4 attention layers out of 56 on the same data.

The pure Mamba models began to fail at looking up a number in a phone book beyond about 500 tokens. The hybrid exceeded the Transformer on all 12 standard tasks the authors evaluated, by 2.65 points on average[13].

Hybrids in practice

The first large hybrid was Jamba from AI21, in March 2024: Mamba layers with one attention layer in every eight[14]. MiniMax, NVIDIA, Qwen and Moonshot followed, each with its own recurrent layer and its own ratio.

MiniMax-Text-01 put one softmax layer after every seven layers of lightning attention, a fast form of linear attention, after finding that lightning attention alone “demonstrates limited retrieval capabilities”[15]. NVIDIA’s Nemotron-H models keep “roughly 8% of the total number of layers” as attention, and the rest are Mamba-2 and MLP layers[16].

Qwen3-Next uses Gated DeltaNet in three layers out of four. According to the Qwen team, at this ratio “the model consistently outperforms any monolithic architecture”[17].

Kimi Linear uses Kimi Delta Attention, the variant with a gate for every channel, also at 3:1. It leaves word order entirely to the recurrent layers: its attention layers use no positional encoding at all[9].

Here is where each model keeps its full attention, read from the config files the models are published with:

The layers of real models, one mark per token-mixing layer, read from their config files. Jamba: 4 of 32 layers are full attention. MiniMax-Text-01: 10 of 80 layers are full attention. Nemotron-H-8B: 4 of 28 layers are full attention. Qwen3-Next-80B-A3B: 12 of 48 layers are full attention. Kimi Linear 48B-A3B: 7 of 27 layers are full attention. Qwen3.5-397B-A17B: 15 of 60 layers are full attention. Kimi K3: 24 of 93 layers are full attention. gpt-oss-20b: 12 of 24 layers are full attention, 12 use a sliding window. Gemma 4 31B: 10 of 60 layers are full attention, 50 use a sliding window. MiniMax-M2: 62 of 62 layers are full attention. MiniMax-M2, the successor of MiniMax-Text-01, uses full attention in all of its layers The layers of real models, one mark per token-mixing layer, read from their config files. Jamba: 4 of 32 layers are full attention. MiniMax-Text-01: 10 of 80 layers are full attention. Nemotron-H-8B: 4 of 28 layers are full attention. Qwen3-Next-80B-A3B: 12 of 48 layers are full attention. Kimi Linear 48B-A3B: 7 of 27 layers are full attention. Qwen3.5-397B-A17B: 15 of 60 layers are full attention. Kimi K3: 24 of 93 layers are full attention. gpt-oss-20b: 12 of 24 layers are full attention, 12 use a sliding window. Gemma 4 31B: 10 of 60 layers are full attention, 50 use a sliding window. MiniMax-M2: 62 of 62 layers are full attention. MiniMax-M2, the successor of MiniMax-Text-01, uses full attention in all of its layers

Drawn to scale, the designs differ more than their headline ratios. Jamba and MiniMax-Text-01 keep one attention layer in eight, Nemotron-H-8B 4 of its 28 token-mixing layers, the Qwen and Kimi models one in four.

Kimi Linear’s released model has 27 layers, so its 3:1 turns into 20 recurrent layers and 7 attention layers, the last one attention.

Kimi K3, 93 layers deep, adds an extra attention layer at the top, “ensuring that the final layer always performs global attention”[18]. All of them spread the attention layers evenly through the depth.

The second group shrinks the cache in another way. gpt-oss and Gemma use softmax attention in every layer, but most of their layers see only a sliding window of recent tokens: 128 in gpt-oss, where windowed and full layers alternate[19], and 1,024 in Gemma 3, with five windowed layers for every full one[20].

Gemma 4 kept the ratio. A window is the opposite trade to a recurrent state: it keeps recent tokens exactly and forgets older ones completely, while a state keeps a blurred summary of everything.

The last row is MiniMax-M2, the successor of MiniMax-Text-01, with full attention in all 62 layers. The team explained why in a blog post.

Small hybrids had matched full attention on the leaderboards, but “the price paid became obvious at a larger scale: the model had clear deficits in complex, multi-hop reasoning tasks”[21].

The Qwen and Kimi hybrids keep more attention than MiniMax-Text-01 did, one layer in four rather than one in eight, and use the delta rule, which lightning attention lacks. Still, the episode is a fair warning that parity on small models and leaderboards doesn’t settle the question.

What the hybrids save is easiest to see in memory. For each model, the size of its cache follows from the config: the number of attention layers, key-value heads and head dimension.

Keys and values one sequence of 128K tokens keeps in bf16, derived from each model's config: Qwen3-32B 32 GiB, MiniMax-M2 31 GiB, Qwen3-8B 18 GiB, MiniMax-Text-01 5 GiB, gpt-oss-20b 3 GiB, Qwen3-Next 3 GiB, Jamba 2 GiB, Kimi Linear 0.98 GiB, Nemotron 3 Nano 0.75 GiB. Qwen3-32B, MiniMax-M2 and Qwen3-8B use full attention in every layer, gpt-oss-20b a sliding window in half of them, the others recurrent layers with full attention in a few Keys and values one sequence of 128K tokens keeps in bf16, derived from each model's config: Qwen3-32B 32 GiB, MiniMax-M2 31 GiB, Qwen3-8B 18 GiB, MiniMax-Text-01 5 GiB, gpt-oss-20b 3 GiB, Qwen3-Next 3 GiB, Jamba 2 GiB, Kimi Linear 0.98 GiB, Nemotron 3 Nano 0.75 GiB. Qwen3-32B, MiniMax-M2 and Qwen3-8B use full attention in every layer, gpt-oss-20b a sliding window in half of them, the others recurrent layers with full attention in a few

For one sequence of 128K tokens, Qwen3-Next keeps 3 GiB of keys and values, against 32 GiB for Qwen3-32B. Not all of the factor of ten comes from the hybrid. Each attention layer of Qwen3-Next keeps two key-value heads of 256 dimensions, where Qwen3-32B keeps eight of 128, which halves the cache per layer.

The rest, a factor of 5.3, is 12 attention layers instead of 64. The recurrent layers add a state of 36 MiB that stays the same at any length.

Kimi Linear starts from multi-head latent attention, introduced in DeepSeek-V2[22], which already caches one compressed vector of 576 numbers per token and layer in place of keys and values for every head, and keeps it in 7 layers of 27.

That is 26% of the cache of an all-MLA model of the same depth, in line with the paper’s “reducing KV cache usage by up to 75%”. MiniMax-M2, at 31 GiB, keeps six times more than its predecessor.

At long context, the cache rather than the weights decides how many conversations fit on a GPU at once, and decoding speed follows from that.

The Kimi Linear authors measured decoding with a million tokens of context and put the result down to memory: the smaller cache “enables larger batches, yielding a 6.3× faster TPOT (1.84ms vs. 11.48ms) than MLA at 1M tokens”, TPOT being the time per output token.

attention in 7 of 27 layers26% of the cachelarger batches6.3× faster decoding at 1M tokens

What it costs on a GPU

The recurrent layer pays off only where the cache grows large, and a single layer on one GPU shows where that starts.

I timed one layer with 16 heads of 128 dimensions, about the size of a Gated DeltaNet layer in Qwen3-Next, in two situations: reading a prompt of \(n\) tokens at once, and writing one new token after \(n\) tokens.

Softmax attention runs through PyTorch’s fused attention kernel in half precision. Gated DeltaNet runs my plain-PyTorch chunked version in single precision, a Python loop over chunks of 64. The comparison is tilted against the recurrent layer on purpose: if it wins anyway, it wins on arithmetic, not on engineering.

Time for one layer with 16 heads of 128 to read a prompt, on a Tesla T4, on logarithmic axes. Softmax attention takes 0.18 ms for 512 tokens and 1,579 ms for 65,536; the chunked Gated DeltaNet in plain PyTorch takes 3.0 ms and 377 ms. Attention is faster up to 8K tokens, the recurrent layer from 16K on Time for one layer with 16 heads of 128 to read a prompt, on a Tesla T4, on logarithmic axes. Softmax attention takes 0.18 ms for 512 tokens and 1,579 ms for 65,536; the chunked Gated DeltaNet in plain PyTorch takes 3.0 ms and 377 ms. Attention is faster up to 8K tokens, the recurrent layer from 16K on

Writing one token after a context of 512 to 65,536 tokens. Left: softmax attention takes 0.08 ms to 3.7 ms, growing with the context; Gated DeltaNet about 0.2 ms at every length. Right: attention reads 4 to 512 MiB of keys and values, the recurrent layer a state of 1 MiB Writing one token after a context of 512 to 65,536 tokens. Left: softmax attention takes 0.08 ms to 3.7 ms, growing with the context; Gated DeltaNet about 0.2 ms at every length. Right: attention reads 4 to 512 MiB of keys and values, the recurrent layer a state of 1 MiB

Reading a prompt, softmax attention is faster up to 8K tokens, which is no surprise: its kernel is one fused program, and mine is a loop. The slopes differ, though.

From 512 to 65,536 tokens, a sequence 128 times longer, the attention layer gets 8,600 times slower and the recurrent one 124 times slower. The curves cross between 8K and 16K tokens, and at 65K the recurrent layer is 4.2 times faster.

Writing a token is where the cache costs most. Attention reads the keys and values of every earlier token, 512 MiB at 65K tokens in this one layer, and its time grows in proportion, from 0.08 ms to 3.7 ms. The recurrent layer updates a state of 1 MiB and takes 0.2 ms at any length; at that size the time is mostly the overhead of launching a few small GPU kernels.

From about 2K tokens on, the recurrent step is faster, and at 65K it is 16 times faster.

On one GPU the recurrent layer overtakes attention from about 2K tokens when writing a token, and between 8K and 16K tokens when reading a prompt.

One layer and one sequence understate the difference in a real server. There, the caches of all the conversations in a batch have to fit in memory at once, so a model whose cache is a quarter or an eighth of the size can serve proportionally more conversations at the same length.

Kimi Linear’s 6.3 times, quoted above, comes from exactly that. The attention layers of a hybrid still pay the full price; three quarters of its layers stop paying it.

When a hybrid is the right choice

A hybrid pays off when contexts are long and there are many of them: agents that carry long transcripts of tool calls, reasoning models that write tens of thousands of tokens, questions over long documents, a server holding hundreds of such conversations at once.

There the cache is most of the memory and reading it is most of the decoding time, and cutting it by a factor of four to ten changes how many users one GPU can serve.

At short contexts the gain is smaller than the headlines suggest. According to Qwen, the decoding throughput of Qwen3-Next is nearly 4 times that of Qwen3-32B at 4K tokens and over 10 times beyond 32K[17].

But Qwen3-Next activates 3 billion parameters per token and Qwen3-32B all 32 billion, so the comparison mixes two savings, the mixture of experts and the hybrid, and the blog doesn’t separate them. The hybrid’s share is the part that grows with the context.

The costs are of three kinds. The first is quality on the tasks attention is best at, precise retrieval and multi-hop reasoning over long inputs: the deficits that sent MiniMax back to full attention.

The second is engineering. Recurrent layers need their own kernels, and serving tricks built around a cache of separate tokens get harder when part of the model keeps a state that has absorbed all of them.

Reusing the cache of a shared prompt needs copies of the state saved at chosen positions, and rolling back a rejected speculative token needs the state from a step earlier.

The third is that the field keeps moving. Qwen3.8-Flash-Next replaced the full attention layers of its hybrid with sparse attention[23], and DeepSeek-V4 combines two kinds of compressed attention instead of adding recurrent layers[24]. Both directions go after the same cache.

A hybrid’s way is to give most layers a memory of fixed size and to keep a few that can still read every token.

References

  1. Katharopoulos, A., Vyas, A., Pappas, N., & Fleuret, F. (2020). Transformers are RNNs: fast autoregressive Transformers with linear attention. ICML 2020. ↩
  2. Arora, S., et al. (2024). Simple linear attention language models balance the recall-throughput tradeoff. ICML 2024. ↩
  3. Widrow, B., & Hoff, M. E. (1960). Adaptive switching circuits. IRE WESCON Convention Record, 4, 96–104. ↩
  4. Schlag, I., Irie, K., & Schmidhuber, J. (2021). Linear Transformers are secretly fast weight programmers. ICML 2021. ↩
  5. Sun, Y., et al. (2023). Retentive network: a successor to Transformer for large language models. arXiv. ↩
  6. Yang, S., Wang, B., Shen, Y., Panda, R., & Kim, Y. (2024). Gated linear attention Transformers with hardware-efficient training. ICML 2024. ↩
  7. Dao, T., & Gu, A. (2024). Transformers are SSMs: generalized models and efficient algorithms through structured state space duality. ICML 2024. ↩
  8. Yang, S., Kautz, J., & Hatamizadeh, A. (2025). Gated delta networks: improving Mamba2 with delta rule. ICLR 2025. ↩
  9. Kimi Team (2025). Kimi Linear: an expressive, efficient attention architecture. arXiv. ↩
  10. Yang, S., Wang, B., Zhang, Y., Shen, Y., & Kim, Y. (2024). Parallelizing linear Transformers with the delta rule over sequence length. NeurIPS 2024. ↩
  11. Arora, S., et al. (2023). Zoology: measuring and improving recall in efficient language models. ICLR 2024. ↩
  12. Jelassi, S., Brandfonbrener, D., Kakade, S. M., & Malach, E. (2024). Repeat after me: Transformers are better than state space models at copying. ICML 2024. ↩
  13. Waleffe, R., et al. (2024). An empirical study of Mamba-based language models. arXiv. ↩
  14. Lieber, O., et al. (2024). Jamba: a hybrid Transformer-Mamba language model. arXiv. ↩
  15. MiniMax (2025). MiniMax-01: scaling foundation models with lightning attention. arXiv. ↩
  16. NVIDIA (2025). Nemotron-H: a family of accurate and efficient hybrid Mamba-Transformer models. arXiv. ↩
  17. Qwen Team (2025). Qwen3-Next: towards ultimate training & inference efficiency. Alibaba Cloud blog. ↩
  18. Kimi Team (2026). Kimi K3: open frontier intelligence. arXiv. ↩
  19. OpenAI (2025). gpt-oss-120b & gpt-oss-20b model card. arXiv. ↩
  20. Gemma Team (2025). Gemma 3 technical report. arXiv. ↩
  21. MiniMax (2025). Why did MiniMax M2 end up as a full attention model? Hugging Face blog. ↩
  22. DeepSeek-AI (2024). DeepSeek-V2: a strong, economical, and efficient mixture-of-experts language model. arXiv. ↩
  23. Qwen Team (2026). Qwen3.8-Flash-Next model card. Hugging Face. ↩
  24. DeepSeek-AI (2026). DeepSeek-V4: towards highly efficient million-token context intelligence. arXiv. ↩
]]> Eva Koroleva It works on my laptop: a reproducible Docker image of a model 2025-12-31T18:00:00+00:00 2026-11-02T00:00:00+00:00 https://blog.korolevalab.com/blog/2026/transformers-from-scratch-09-docker Transformers: from scratch to production · lesson 9 of 10 · all lessons · run on Kaggle · open in Colab · code on GitHub

On 16 June 2024 NumPy 2.0.0 came out, and code that nobody had touched started to fail. A model served with PyTorch 2.3.1 printed this on import, and stopped [1]:

A module that was compiled using NumPy 1.x cannot be run in
NumPy 2.0.0 as it may crash. To support both 1.x and 2.x
versions of NumPy, modules must be compiled with NumPy 2.0.
Some module may need to rebuild instead e.g. with 'pybind11>=2.12'.

Nothing in those projects had changed. Their requirements said numpy, or numpy>=1.24, and on that Sunday the newest NumPy that matched became a different program.

NumPy had warned about it in its release notes: packages “built against a NumPy 1.xx release will not work with NumPy 2.0” [1]. PyTorch caught up five weeks later, in 2.4.0. Everyone who reinstalled in between learned what their model was really made of.

It is made of more than weights. When perevod translates a sentence, its 38 MB of weights are one part of what runs: around them are the package’s code, three dozen Python packages, Python itself, the system libraries that Python and PyTorch load, and a processor.

Change any of them, and you have a different program, which may or may not give the same translations, or start at all.

What runs when perevod translates, as a stack: the weights, perevod's code, the Python packages, Python and Debian's libraries, the build itself, and the processor. Next to each, what pins it in perevod 0.3.1: the model's commit, a git tag, a lock file with hashes, a base image pinned by digest, a pinned builder and a fixed time; the processor is the one part no image can pin What runs when perevod translates, as a stack: the weights, perevod's code, the Python packages, Python and Debian's libraries, the build itself, and the processor. Next to each, what pins it in perevod 0.3.1: the model's commit, a git tag, a lock file with hashes, a base image pinned by digest, a pinned builder and a fixed time; the processor is the one part no image can pin

This lesson pins every layer of that stack except the last, and puts the result in one Docker image that gives the same bytes each time it’s built, so that anyone can check it was built from the code they can read.

Then it takes the image apart without Docker, runs the model from it, and finally ships the same model a second way: to the browser, where a free page translates on the visitor’s own device.

What you’ll do

  1. Watch the same requirements turn into three different programs on three different days.
  2. Freeze them in a lock file with hashes, and see pip refuse a file that doesn’t match.
  3. Take a container image apart by hand: a registry, manifests, layers and digests.
  4. Find out why two builds of one Dockerfile differ, and remove each cause.
  5. Write perevod’s Dockerfile, and a GitHub Actions job that builds the image twice and publishes it.
  6. Pull the published image without Docker, run the model from it, and check its weights.
  7. Export the model to ONNX and run it in a browser, on a free Hugging Face Space.

What you need: a CPU, as in lessons 7 and 8. Not Docker: Kaggle and Colab don’t have it, so the notebook does by hand what Docker does, which turns out to be the best way to see what an image is.

It downloads the published image (about 365 MB) and runs in a few minutes. On your own machine with Docker, the lesson also shows the commands that build and run the image.

Setup

One small helper package, tfs, draws the course’s plots and token chips and knows where the notebook runs (Kaggle, Colab or your own machine). Everything the lesson teaches is written out in the notebook itself.

import importlib.util
import subprocess
import sys

TFS = "git+https://github.com/xmarva/transformers-from-scratch@main"
if importlib.util.find_spec("tfs") is None:
    subprocess.check_call([sys.executable, "-m", "pip", "install", "-q", TFS])

from tfs import env
from tfs.viz import table, token_chips, use_course_style

env.quiet()  # no progress bars or environment warnings in the outputs
SEED = 42
env.seed_everything(SEED)
COLORS = use_course_style()
WORK_DIR = env.work_dir()

import gzip
import hashlib
import io
import json
import math
import os
import py_compile
import shutil
import tarfile
import time
from pathlib import Path

import numpy as np
import pandas as pd
import requests
import torch
from huggingface_hub import hf_hub_download

print("Python", sys.version.split()[0], "on", os.uname().machine)

Python 3.13.15 on x86_64

1. One requirements list, three programs

Start from perevod as lesson 8 released it, version 0.2.0: the release’s tag, and only the perevod/ folder of the course repository.

REPO = WORK_DIR / "course"
shutil.rmtree(REPO, ignore_errors=True)
URL = "https://github.com/xmarva/transformers-from-scratch"
QUIET = ["-c", "advice.detachedHead=false"]
subprocess.run(["git", *QUIET, "clone", "-q", "--depth", "1", "--branch", "perevod-v0.2.0",
                "--filter=blob:none", "--sparse", URL, str(REPO)], check=True)
subprocess.run(["git", "sparse-checkout", "set", "perevod"], cwd=REPO, check=True)
os.chdir(REPO)


def write(path, text):
    Path(path).parent.mkdir(parents=True, exist_ok=True)
    Path(path).write_text(text.lstrip("\n"), encoding="utf-8")


def replace(path, old, new):
    text = Path(path).read_text(encoding="utf-8")
    assert text.count(old) == 1, old
    Path(path).write_text(text.replace(old, new), encoding="utf-8")


print(Path("perevod/pyproject.toml").read_text().split("[project.scripts]")[0])

[project]
name = "perevod"
dynamic = ["version"]  # read from src/perevod/__init__.py, see [tool.hatch.version]
description = "English-to-Russian translation with a small Transformer built from scratch"
readme = "README.md"
license = "MIT"
requires-python = ">=3.10"
dependencies = [
    "torch>=2.1",
    "tokenizers>=0.15",
    "huggingface_hub>=0.20",
    "safetensors>=0.4",
]

[project.optional-dependencies]
serve = ["fastapi>=0.110", "uvicorn>=0.29"]
demo = ["fastapi>=0.110", "uvicorn>=0.29", "gradio>=4.40"]
test = ["pytest>=8", "ruff>=0.6", "sacrebleu>=2.4", "pandas>=2.0", "pyarrow>=14",
        "fastapi>=0.110", "httpx>=0.27"]

Its dependencies are ranges: torch>=2.1 means “any PyTorch from 2.1 on”. That’s right for a package other people install. perevod has to live next to whatever else they have, and the Python Packaging User Guide calls pinning exact versions there “overly-restrictive” [2]. But a range is a question, and pip answers it on the day you install.

You can ask it on other days. uv, a fast replacement for pip written in Rust, can resolve requirements as if it were an earlier date: --exclude-newer makes it ignore every file uploaded after it.

Here is what perevod’s requirements, with the server extra, resolved to a year before this lesson, half a year before, and on the day perevod’s lock file was made:

!{sys.executable} -m pip install -q uv==0.12.24
UV = [sys.executable, "-m", "uv"]


def resolve(day):
    command = [*UV, "pip", "compile", "perevod/pyproject.toml", "--extra", "serve",
               "--python-version", "3.12", "--python-platform", "x86_64-manylinux_2_28",
               "--torch-backend", "cpu", "--exclude-newer", day,
               "--no-header", "--no-annotate", "--quiet"]
    result = subprocess.run(command, capture_output=True, text=True, check=True)
    versions = {}
    for line in result.stdout.splitlines():
        name, version = line.split("==")
        versions[name] = version
    return versions


DAYS = ["2025-10-27T00:00:00Z", "2026-04-27T00:00:00Z", "2026-10-27T00:00:00Z"]
resolved = {}
for day in DAYS:
    resolved[day[:10]] = resolve(day)

DIRECT = ["torch", "tokenizers", "huggingface-hub", "safetensors", "fastapi", "uvicorn"]
rows = []
for name in DIRECT:
    row = {"package": name}
    for day, versions in resolved.items():
        row[day] = versions[name]
    rows.append(row)
display(table(pd.DataFrame(rows)))

first, last = resolved["2025-10-27"], resolved["2026-10-27"]
changed = []
for name in first:
    if name in last and first[name] != last[name]:
        changed.append(name)
print(len(first), "packages a year ago,", len(last), "on the release day")
print(len(changed), "packages installed on both days, at different versions")
print("gone:", sorted(set(first) - set(last)))
print("new: ", sorted(set(last) - set(first)))
package2025-10-272026-04-272026-10-27
torch2.9.0+cpu2.10.0+cpu2.14.1+cpu
tokenizers0.22.10.22.20.23.3
huggingface-hub0.36.01.12.02.3.0
safetensors0.6.20.7.00.8.0
fastapi0.120.00.136.10.162.0
uvicorn0.38.00.46.00.54.0
34 packages a year ago, 33 on the release day
24 packages installed on both days, at different versions
gone: ['certifi', 'charset-normalizer', 'requests', 'sniffio', 'urllib3']
new:  ['httpcore2', 'httpx2', 'opentelemetry-api', 'truststore']

Same pyproject.toml, three programs: of the packages installed both a year ago and on the release day, 24 are at different versions.

PyTorch moves by several releases, and each of them changes kernels, defaults and warnings. huggingface-hub, which downloads the model, jumps two major versions.

And these are only the packages perevod names: behind them, nearly every package it pulls in changed version too, and some were replaced outright, an HTTP library by another. None of this is a bug. It’s what a range asks for.

For a library, that’s the deal. For a service, it’s a gamble: the server you test today and the one that restarts next month install different code. Donald Stufft, who maintained pip and PyPI for years, drew the line in 2013[3]:

While a library tends to want to have wide open ended version specifiers an application wants very specific dependencies.

Donald Stufft

perevod is both. It keeps the ranges in pyproject.toml for people who install it as a library, and the service is deployed from a second file that pins every package: a lock file.

2. A lock file, with hashes

First, a change that the image revealed. On a laptop, some other package has always installed NumPy; in an image that holds only what perevod asks for, it’s missing, and PyTorch, which tries to import NumPy when it loads, prints Failed to initialize NumPy: No module named 'numpy'. perevod never calls NumPy, so nothing breaks, but a server that starts with a warning trains people to ignore warnings. perevod now asks for it:

replace("perevod/pyproject.toml", '    "safetensors>=0.4",\n',
        '    "safetensors>=0.4",\n'
        '    "numpy>=1.24",  # perevod doesn\'t use it, '
        'but PyTorch warns on import without it\n')

Now the lock. uv pip compile resolves the requirements once and writes the answer down, one exact version per package. Three options matter here:

  • --universal: one file for every platform. Where the answer differs, a line carries a marker, such as ; sys_platform == 'win32', and pip skips the lines that don’t apply.
  • --torch-backend cpu: PyTorch’s CPU-only build from PyTorch’s own index. The default Linux wheel on PyPI is built with CUDA and pulls in NVIDIA’s libraries: for torch 2.14.1 that’s about 3 GB of downloads, against 196 MB for the CPU build. A translation server without a GPU needs none of it.
  • --generate-hashes: next to each version, the sha256 of every file that version may be installed from.

And one more for this lesson: --exclude-newer with the day perevod 0.3’s lock was made, so that the notebook resolves the same answer even if you run it a year from now.

LOCKED_ON = "2026-10-27T12:00:00Z"  # the day perevod 0.3's lock was made
COMPILE = [*UV, "pip", "compile", "--universal", "--python-version", "3.12",
           "--generate-hashes", "--exclude-newer", LOCKED_ON, "--quiet"]
subprocess.run([*COMPILE, "pyproject.toml", "--extra", "serve", "--torch-backend", "cpu",
                "-o", "requirements.lock"], cwd="perevod", check=True)

lock = Path("perevod/requirements.lock").read_text()
pins = []
for line in lock.splitlines():
    if "==" in line and not line.startswith((" ", "#")):
        pins.append(line.rstrip(" \\"))
print(len(pins), "pinned lines,", lock.count("--hash=sha256:"), "hashes")
for pin in pins:
    if pin.startswith(("torch==", "numpy==", "huggingface-hub==")):
        print(pin)

37 pinned lines, 582 hashes
huggingface-hub==2.3.0
numpy==2.5.4
torch==2.14.1 ; sys_platform == 'darwin'
torch==2.14.1+cpu ; sys_platform != 'darwin'

PyTorch gets two lines. On macOS, PyPI’s wheel is already CPU-only and the PyTorch index has no +cpu build, so the lock takes the plain one there and the +cpu build everywhere else.

The hashes are what make the lock more than a list of versions. A version is a name: whoever controls the index, a mirror, or a proxy between you and it, decides which bytes stand behind it.

A hash is the bytes themselves. With a hash on every line, pip switches to hash-checking mode, in which “Hashes are required for all requirements” and every downloaded file must match one of them [4].

Try it on one small package from the lock, h11, an HTTP/1.1 library the server uses. First as locked:

lines = lock.splitlines()
start = 0
for k, line in enumerate(lines):
    if line.startswith("h11=="):
        start = k
entry = [lines[start]]
for line in lines[start + 1:]:
    if not line.startswith("    --hash"):
        break
    entry.append(line)
good = "\n".join(entry).rstrip(" \\") + "\n"
print(good)

PIP_DOWNLOAD = [sys.executable, "-m", "pip", "download", "--no-deps", "--require-hashes",
                "--no-color", "--disable-pip-version-check"]
shutil.rmtree(WORK_DIR / "downloads", ignore_errors=True)
Path(WORK_DIR / "good.txt").write_text(good)
subprocess.run([*PIP_DOWNLOAD, "-q", "-d", str(WORK_DIR / "downloads"),
                "-r", str(WORK_DIR / "good.txt")], check=True)
for path in sorted((WORK_DIR / "downloads").glob("h11-*")):
    print(path.name, "sha256:" + hashlib.sha256(path.read_bytes()).hexdigest())

h11==0.16.0 \
    --hash=sha256:4e35b956cf45792e4caa5885e69fba00bdbc6ffafbfa020300e549b208ee5ff1 \
    --hash=sha256:63cf8bbe7522de3bf65932fda1d9c2772064ffb3dae62d55932da54b31cb6c86

h11-0.16.0-py3-none-any.whl sha256:63cf8bbe7522de3bf65932fda1d9c2772064ffb3dae62d55932da54b31cb6c86

The file pip downloaded has one of the locked hashes. Now change the last hex digit of each hash, which is what pip would see if someone swapped the file behind the version:

bad_lines = []
for line in good.splitlines():
    if "--hash=sha256:" in line:
        digest = line.split("sha256:")[1][:64]
        last = "1" if digest[-1] == "0" else "0"
        line = line.replace(digest, digest[:-1] + last)
    bad_lines.append(line)
Path(WORK_DIR / "bad.txt").write_text("\n".join(bad_lines) + "\n")
shutil.rmtree(WORK_DIR / "tampered", ignore_errors=True)
result = subprocess.run([*PIP_DOWNLOAD, "-d", str(WORK_DIR / "tampered"),
                         "-r", str(WORK_DIR / "bad.txt")], capture_output=True, text=True)
print("pip exit code:", result.returncode)
print(result.stderr.strip().replace(str(WORK_DIR) + "/", ""))
print("files downloaded:", list((WORK_DIR / "tampered").glob("*")))

pip exit code: 1
ERROR: THESE PACKAGES DO NOT MATCH THE HASHES FROM THE REQUIREMENTS FILE. If you have updated the package versions, please update the hashes. Otherwise, examine the package contents carefully; someone may have tampered with them.
    h11==0.16.0 from https://files.pythonhosted.org/packages/04/4b/29cac41a4d98d144bf5f6d33995617b185d14b22401f75ca86f384e87ff1/h11-0.16.0-py3-none-any.whl (from -r bad.txt (line 1)):
        Expected sha256 4e35b956cf45792e4caa5885e69fba00bdbc6ffafbfa020300e549b208ee5ff0
        Expected     or 63cf8bbe7522de3bf65932fda1d9c2772064ffb3dae62d55932da54b31cb6c80
             Got        63cf8bbe7522de3bf65932fda1d9c2772064ffb3dae62d55932da54b31cb6c86
files downloaded: []

pip refuses, names the package, and shows the hashes the lock expects next to the one it got. Nothing is kept, let alone installed. The lock now pins the bytes of every file, not just a version number.

One more program runs during the install, and nobody asked for it by name: the build backend. To install perevod from source, pip first builds a wheel with hatchling, the backend pyproject.toml names, and by default it builds in a fresh environment with whatever hatchling is newest that day.

A newer hatchling writes slightly different metadata into the wheel, so it gets a lock of its own, and the Dockerfile will build without that fresh environment:

write("perevod/build.in", "hatchling\n")
subprocess.run([*COMPILE, "build.in", "-o", "build.lock"], cwd="perevod", check=True)
for line in Path("perevod/build.lock").read_text().splitlines():
    if "==" in line and not line.startswith((" ", "#")):
        print(line.rstrip(" \\"))

hatchling==1.32.4
packaging==26.3
pathspec==1.1.1
pluggy==1.7.0
tomlkit==0.15.2
trove-classifiers==2026.9.21.13

3. What an image is

A Docker image sounds like a disk image, one big file. It isn’t: it’s a few small JSON documents and a stack of tar archives, stored in a registry under names that are their own sha256 hashes. All of it comes over plain HTTP, and fetching it once by hand explains most of the rest of this lesson.

Registries speak one protocol, the OCI distribution spec: Docker Hub, GitHub’s registry and the others. To pull, a client asks /v2/<name>/manifests/<tag> for a manifest, then downloads what the manifest lists from /v2/<name>/blobs/<digest> (OCI distribution spec).

Public images need no account, only a short-lived anonymous token. Here are the four requests, as functions:

INDEX = "application/vnd.oci.image.index.v1+json"
MANIFEST = "application/vnd.oci.image.manifest.v1+json"
DOCKER_LIST = "application/vnd.docker.distribution.manifest.list.v2+json"
DOCKER_MANIFEST = "application/vnd.docker.distribution.manifest.v2+json"


def sha256(data):
    return "sha256:" + hashlib.sha256(data).hexdigest()


def pull_token(registry, repository):
    scope = f"repository:{repository}:pull"
    if registry == "registry-1.docker.io":
        params = {"service": "registry.docker.io", "scope": scope}
        response = requests.get("https://auth.docker.io/token", params=params)
    else:
        response = requests.get(f"https://{registry}/token", params={"scope": scope})
    return response.json()["token"]


def get_manifest(registry, repository, reference, token):
    accept = ", ".join([INDEX, MANIFEST, DOCKER_LIST, DOCKER_MANIFEST])
    headers = {"Authorization": f"Bearer {token}", "Accept": accept}
    url = f"https://{registry}/v2/{repository}/manifests/{reference}"
    response = requests.get(url, headers=headers)
    response.raise_for_status()
    return response


def get_blob(registry, repository, digest, token):
    headers = {"Authorization": f"Bearer {token}"}
    url = f"https://{registry}/v2/{repository}/blobs/{digest}"
    response = requests.get(url, headers=headers)
    response.raise_for_status()
    assert sha256(response.content) == digest  # the name of a blob is the hash of its bytes
    return response.content

perevod’s image is built on python:3.12-slim, the official Python image on Debian. Ask Docker Hub for it:

HUB = "registry-1.docker.io"
hub_token = pull_token(HUB, "library/python")
response = get_manifest(HUB, "library/python", "3.12-slim", hub_token)
print("Content-Type:         ", response.headers["Content-Type"])
print("Docker-Content-Digest:", response.headers["Docker-Content-Digest"])
print("sha256 of the body:   ", sha256(response.content))

index = response.json()
rows = []
attestations = 0
for entry in index["manifests"]:
    platform = entry["platform"]
    if platform["architecture"] == "unknown":
        attestations += 1
        continue
    name = platform["os"] + "/" + platform["architecture"]
    if "variant" in platform:
        name += "/" + platform["variant"]
    rows.append({"platform": name, "manifest": entry["digest"][:19] + "…",
                 "bytes": entry["size"]})
display(table(pd.DataFrame(rows)))
print("and", attestations, "attestation manifests, one per image")

Content-Type:          application/vnd.oci.image.index.v1+json
Docker-Content-Digest: sha256:aa943ce411423c7582288f18707ff66eddcbaeb510502d747a72aeb34f30d159
sha256 of the body:    sha256:aa943ce411423c7582288f18707ff66eddcbaeb510502d747a72aeb34f30d159
platformmanifestbytes
linux/amd64sha256:1148dbb01461…1,746
linux/arm/v5sha256:e353f56390bc…1,748
linux/arm/v7sha256:e3f11fed4d77…1,748
linux/arm64/v8sha256:b224e2fe306f…1,748
linux/386sha256:0f83099f3138…1,745
linux/ppc64lesha256:20d5360eb0a6…1,748
linux/riscv64sha256:2628eda9aec1…1,748
linux/s390xsha256:c806450897d2…1,746
and 8 attestation manifests, one per image

The digest the registry reports is the sha256 of the bytes it sent, and the OCI image spec makes that the content’s name: a digest “acts as a content identifier, enabling content addressability” (OCI image spec: descriptors). Change one byte, and it’s a different object with a different name.

What came back is an image index: one image per processor architecture, from x86-64 to IBM’s mainframes, plus as many attestations, records of how each image was built. docker pull takes the entry for your machine. Follow the one for x86-64 down to its config and its layers:

for entry in index["manifests"]:
    if entry["platform"]["architecture"] == "amd64":
        amd64 = entry
manifest = get_manifest(HUB, "library/python", amd64["digest"], hub_token).json()
config_digest = manifest["config"]["digest"]
config = json.loads(get_blob(HUB, "library/python", config_digest, hub_token))
print("built:", config["created"])
for item in config["config"]["Env"]:
    if item.startswith("PYTHON_VERSION"):
        print(item)

steps = []
for step in config["history"]:
    if not step.get("empty_layer"):
        steps.append(step["created_by"])
rows = []
for layer, step in zip(manifest["layers"], steps, strict=True):
    command = " ".join(step.split())
    rows.append({"layer": layer["digest"][:19] + "…", "MB": layer["size"] / 1e6,
                 "made by": command[:58]})
display(table(pd.DataFrame(rows), digits=1))

built: 2026-10-28T02:03:39.583192962Z
PYTHON_VERSION=3.12.16
layerMBmade by
sha256:6fc242d221a6…29.9# debian.sh --arch 'amd64' out/ 'trixie' '@1793059200'
sha256:3eb771151506…1.3RUN /bin/sh -c set -eux; apt-get update; apt-get install -
sha256:f574edac59d9…12.1RUN /bin/sh -c set -eux; savedAptMark="$(apt-mark showmanu
sha256:acca7542dfc7…0.0RUN /bin/sh -c set -eux; for src in idle3 pip3 pydoc3 pyth

That’s the whole structure:

The structure of an image, top to bottom. A tag points to an image index. The index lists one manifest per platform by digest. A manifest lists a config and the layers by digest. The config holds the environment, the command and the hash of each unpacked layer; each layer is a gzipped tar archive of the files one Dockerfile step changed. Each level names the next by its sha256, so the digest at the top fixes every byte below it The structure of an image, top to bottom. A tag points to an image index. The index lists one manifest per platform by digest. A manifest lists a config and the layers by digest. The config holds the environment, the command and the hash of each unpacked layer; each layer is a gzipped tar archive of the files one Dockerfile step changed. Each level names the next by its sha256, so the digest at the top fixes every byte below it

A layer is a gzipped tar archive of the files one step of a Dockerfile added or changed. The config is a JSON document with what Docker needs to start the image: environment variables, the user, the command, and the hash of each unpacked layer.

The manifest lists the config and the layers by digest, and the index lists the manifests. Every level names the next one by its hash, so the digest at the top fixes every byte below it. FROM python:3.12-slim@sha256:… pins the whole base image down to the last file.

Trap

FROM python:3.12-slim pins nothing. A tag is a pointer, and Docker says so plainly: “Image tags are mutable, meaning a publisher can update a tag to point to a new image” (Docker: build best practices).

The official images are rebuilt whenever Debian or Python ships a fix; the maintainers “strive to publish updated images at least monthly for Debian”[5]. Compare today’s tag with the digest perevod’s Dockerfile pins, which was the tag’s target the day the Dockerfile was written:

PINNED = "sha256:12316debbfb86890f6b2d825820612741ba527ad5a44e3902759141d572bfe8f"
today = response.headers["Docker-Content-Digest"]
if today == PINNED:
    print("python:3.12-slim still points to the image perevod pins")
else:
    print("python:3.12-slim has moved: today it points to", today[:19] + "…")
pinned = get_manifest(HUB, "library/python", PINNED, hub_token)
print("the pinned image is still served by its digest:", sha256(pinned.content) == PINNED)

python:3.12-slim has moved: today it points to sha256:aa943ce41142…
the pinned image is still served by its digest: True

4. Why two builds of one Dockerfile differ

The definition the reproducible-builds project uses is strict: “A build is reproducible if given the same source code, build environment and build instructions, any party can recreate bit-by-bit identical copies of all specified artifacts”[6].

For an image, bit by bit means the same digest at the top. It’s worth the trouble for one reason: then anyone can check that the image was built from the source they can read, by building it and comparing one hash.

Docker doesn’t give you that by default. Build one Dockerfile twice without the cache, and you get two digests.

This lesson’s image did exactly that on its first two builds, and the differences were in places that have nothing to do with the program. Two of them you can reproduce right here.

The first is time. A layer is a tar archive, and tar stores each file’s modification time. Pack perevod’s source as a layer, then “save” the same files a minute later, as a fresh checkout in CI would:

commit_time = subprocess.run(["git", "log", "-1", "--format=%ct"], capture_output=True,
                             text=True, check=True)
EPOCH = int(commit_time.stdout)  # the time of the commit the notebook started from


def layer_digest(folder, normalize):
    buffer = io.BytesIO()
    with tarfile.open(fileobj=buffer, mode="w") as tar:
        for path in sorted(Path(folder).rglob("*")):
            info = tar.gettarinfo(str(path), arcname=str(path.relative_to(folder)))
            if normalize:
                info.mtime = min(info.mtime, EPOCH)  # no time later than the commit
                info.uid = 0
                info.gid = 0
                info.uname = ""
                info.gname = ""
            if path.is_file():
                with open(path, "rb") as file:
                    tar.addfile(info, file)
            else:
                tar.addfile(info)
    return sha256(gzip.compress(buffer.getvalue(), mtime=0))  # gzip stores a time too


def touch(path):
    stat = path.stat()
    os.utime(path, (stat.st_atime, stat.st_mtime + 60))


SOURCE = Path("perevod/src")
as_is = layer_digest(SOURCE, normalize=False)
normalized = layer_digest(SOURCE, normalize=True)
for path in SOURCE.rglob("*.py"):
    touch(path)
print("as on disk:", as_is[:23], "->", layer_digest(SOURCE, normalize=False)[:23])
print("normalized:", normalized[:23], "->", layer_digest(SOURCE, normalize=True)[:23])

as on disk: sha256:9e9b126813cbaec2 -> sha256:e673f23034cdb19a
normalized: sha256:c25edb033761b87c -> sha256:c25edb033761b87c

Same bytes in every file, a new digest. Clamping every time to the commit’s time, and dropping the owner’s name, fixes it.

BuildKit, the engine behind docker build, does exactly this when asked: it takes the time from the SOURCE_DATE_EPOCH variable, a convention of the reproducible-builds project, and with the exporter option rewrite-timestamp=true it clamps the time of every file in every layer to it [7].

The second cause hides inside the files. pip compiles every .py file it installs to bytecode, a .pyc, and a .pyc starts with a 16-byte header that tells Python whether the cache is still valid. By default the header holds the source file’s modification time:

module = Path("perevod/src/perevod/config.py")
target = WORK_DIR / "config.pyc"


def compiled():
    py_compile.compile(str(module), cfile=str(target))
    return target.read_bytes()


def header(pyc):
    flags = int.from_bytes(pyc[4:8], "little")
    if flags == 0:
        mtime = int.from_bytes(pyc[8:12], "little")
        size = int.from_bytes(pyc[12:16], "little")
        return f"flags {flags}, source modified at {mtime}, source size {size}"
    return f"flags {flags}, source hash {pyc[8:16].hex()}"


first = compiled()
touch(module)
second = compiled()
print(header(first))
print(header(second))
print("same bytes:", first == second)

flags 0, source modified at 1793220667, source size 1856
flags 0, source modified at 1793220727, source size 1856
same bytes: False

Even if the archive’s times are clamped, the time inside the file survives, and pip installs thousands of .pyc files. PEP 552 added a second kind of header, with a hash of the source instead of its time, and py_compile, which pip uses through compileall, switches to it by itself when SOURCE_DATE_EPOCH is set [8]:

os.environ["SOURCE_DATE_EPOCH"] = str(EPOCH)
first = compiled()
touch(module)
second = compiled()
del os.environ["SOURCE_DATE_EPOCH"]
print(header(first))
print(header(second))
print("same bytes:", first == second)

flags 3, source hash 6e4327a19146f24e
flags 3, source hash 6e4327a19146f24e
same bytes: True

Flags 3 mean “checked hash”: Python compares the stored hash with the source’s when it imports the module. The bytes no longer depend on when the file was saved.

There’s a catch in the Dockerfile, though: BuildKit hands SOURCE_DATE_EPOCH to the build, but a RUN step only sees a build argument that the Dockerfile declares with ARG. Without that one line, pip in the image never sees the variable.

The 16-byte header of a .pyc file, compiled twice from the same source saved a minute apart. Timestamp mode: magic number, flags 0, the source's modification time and size; the time differs, so the files differ. Checked-hash mode, with SOURCE_DATE_EPOCH set: magic number, flags 3, an 8-byte hash of the source; the two files are identical The 16-byte header of a .pyc file, compiled twice from the same source saved a minute apart. Timestamp mode: magic number, flags 0, the source's modification time and size; the time differs, so the files differ. Checked-hash mode, with SOURCE_DATE_EPOCH set: magic number, flags 3, an 8-byte hash of the source; the two files are identical

The rest of the list, from the first naive builds to two identical ones:

What differed Why The fix
file times in every layer each file gets the time it was written SOURCE_DATE_EPOCH and rewrite-timestamp=true
the dates in the image’s config the clock at build time SOURCE_DATE_EPOCH
every .pyc that pip wrote the header holds the source’s time ARG SOURCE_DATE_EPOCH, so pip sees it
a log file in the model’s layer hf-xet names its log by the time delete it in the same RUN
the base image, next month the tag moves FROM python:3.12-slim@sha256:…
the packages, next month new releases the lock file, with hashes
the wheel’s metadata pip fetches the newest hatchling build.lock and --no-build-isolation
the archives, and the parser each BuildKit version writes and parses its own way BuildKit’s version pinned in CI
attestations they record when and where a build ran --provenance=false --sbom=false

The order of the entries in a tar archive can differ too, in general; here BuildKit’s builds agreed on it. And one cause this table missed: perevod 0.3.0 shipped with it, and section 7 shows it.

5. perevod’s Dockerfile

Before the image, two small changes in the package. An image runs somewhere you can’t see, so it should be able to say what it is: /health now reports the model’s commit along with the package’s version. The translator remembers where its weights came from; a snapshot downloaded from the Hub lives in a folder named by its commit:

TRANSLATOR = "perevod/src/perevod/translator.py"
replace(TRANSLATOR, "        self.generation = generation\n",
        "        self.generation = generation\n"
        "        self.source = None  "
        "# where the weights came from; from_pretrained fills it in\n")
replace(TRANSLATOR,
        "        if not path.is_dir():\n"
        "            path = Path(snapshot_download(model, revision=revision, "
        "allow_patterns=FILES))\n",
        "        if path.is_dir():\n"
        "            source = str(path)\n"
        "        else:\n"
        "            path = Path(snapshot_download(model, revision=revision, "
        "allow_patterns=FILES))\n"
        '            source = f"{model}@{path.name}"  '
        "# a snapshot's folder is named by its commit\n")
replace(TRANSLATOR,
        "        return cls(network.to(device), tokenizer, "
        "GenerationConfig.from_dict(settings))\n",
        "        generation = GenerationConfig.from_dict(settings)\n"
        "        translator = cls(network.to(device), tokenizer, generation)\n"
        "        translator.source = source\n"
        "        return translator\n")

The server returns it, and its test checks it:

replace("perevod/src/perevod/server.py",
        "    async def health():\n"
        '        return {"status": "ok", "version": __version__}\n',
        "    async def health():  "
        "# what is running: the package's version and the model's commit\n"
        "        model = app.state.batcher.translator.source\n"
        '        return {"status": "ok", "version": __version__, "model": model}\n')
replace("perevod/tests/test_server.py",
        "def test_the_health_check(client):\n"
        '    assert client.get("/health").json()["status"] == "ok"\n',
        "def test_the_health_check_names_the_model(client, tiny_dir):\n"
        '    health = client.get("/health").json()\n'
        '    assert health["status"] == "ok"\n'
        '    assert health["model"] == str(tiny_dir)\n')
replace("perevod/src/perevod/__init__.py", '__version__ = "0.2.0"', '__version__ = "0.3.1"')

Now the Dockerfile, top to bottom.

FROM takes the Python image by digest, as section 3 showed. The Dockerfile’s parser is part of BuildKit, so pinning BuildKit’s version in CI pins it too. (A # syntax= line on top can name a newer parser, as an image of its own; then pin that by digest as well.)

ARG SOURCE_DATE_EPOCH is the line from section 4 that lets pip see the build’s time. A user with uid 1000 instead of root: a process that has no reason to write to the system shouldn’t be able to.

The dependencies are installed before the code, and exactly as locked: --require-hashes turns on the checks, --no-deps makes pip install only what the lock lists, and the lock already lists everything. The order matters for speed, too.

Docker caches each step and reruns a step only when something it depends on changes, so the slow step that installs PyTorch, which changes least often, comes first, and editing perevod’s code rebuilds only the cheap steps after it.

The package is installed with --no-build-isolation, so that pip builds it with the locked hatchling from the previous step. Then the model, in a folder of its own, at the commit perevod pins.

The download leaves more than the model behind in ~/.cache, and all of that goes in the same step, because a file deleted in a later step stays in the earlier layer forever.

Finally HF_HUB_OFFLINE=1: in the running container, “no HTTP calls will be made to the Hugging Face Hub” (huggingface_hub: environment variables), so the image runs where there is no network, and never quietly downloads anything else. The labels say where the image came from; GitHub uses org.opencontainers.image.source to link the image to its repository.

write("perevod/Dockerfile", """
# perevod's server with its model inside: one commit builds one image, bit for bit.
# The base image by digest; the Dockerfile's parser comes with BuildKit, pinned in CI.
FROM \\
    python:3.12-slim@sha256:12316debbfb86890f6b2d825820612741ba527ad5a44e3902759141d572bfe8f

# a fixed time instead of the clock: with it pip writes .pyc files that don't record when
ARG SOURCE_DATE_EPOCH
ENV PIP_NO_CACHE_DIR=1 PIP_DISABLE_PIP_VERSION_CHECK=1 PYTHONUNBUFFERED=1
RUN useradd --uid 1000 --create-home perevod
WORKDIR /home/perevod

# dependencies first, exactly as locked: they change least often, so this layer is cached
COPY requirements.lock build.lock ./
RUN pip install --require-hashes --no-deps -r requirements.lock -r build.lock \\
        --extra-index-url https://download.pytorch.org/whl/cpu

# the package itself, built with the locked build backend
COPY pyproject.toml README.md ./src/
COPY src ./src/src
RUN pip install --no-deps --no-build-isolation ./src && rm -rf ./src

# the model, at the commit this version of perevod was tested with, in a folder of its own;
# the rest of what the download leaves in ~/.cache goes: hf-xet's log is named by the time,
# and huggingface_hub caches a list of AI agents that the Hub can change at any moment
ENV HF_HUB_CACHE=/home/perevod/models
USER perevod
RUN python -c "from perevod import Translator; Translator.from_pretrained()" \\
    && rm -rf /home/perevod/.cache
ENV HF_HUB_OFFLINE=1

ARG VERSION=dev REVISION=unknown
LABEL org.opencontainers.image.source=https://github.com/xmarva/transformers-from-scratch \\
      org.opencontainers.image.version=$VERSION \\
      org.opencontainers.image.revision=$REVISION
EXPOSE 8000
CMD ["perevod-server", "--host", "0.0.0.0", "--port", "8000"]
""")
assert PINNED in Path("perevod/Dockerfile").read_text()

The build context, the folder Docker sends to the builder, is perevod/. A .dockerignore keeps out what the image doesn’t need, and what would change its bytes without changing the program, like Python’s caches:

write("perevod/.dockerignore", """
# only what the image needs: the code, its metadata and the lock files
tests
**/__pycache__
.pytest_cache
.ruff_cache
dist
build
*.egg-info
""")

On a machine with Docker, this builds the image and runs it, with the same options CI uses below (Docker Desktop and OrbStack ship with BuildKit):

cd perevod
export SOURCE_DATE_EPOCH=$(git log -1 --format=%ct)
docker buildx build . --build-arg SOURCE_DATE_EPOCH --provenance=false --sbom=false \
    --output type=docker,name=perevod:dev,rewrite-timestamp=true
docker run --rm -p 8000:8000 perevod:dev

And without a network at all, from the command line of lesson 7:

docker run --rm --network none perevod:dev perevod "Tom is afraid of the dark."

One trap when you limit the container’s CPUs. docker run --cpus 2 is a quota on CPU time, not a set of processors, and PyTorch doesn’t read quotas: it starts as many threads as the host has physical cores, and they fight over the two CPUs’ worth of time. Tell it explicitly, with -e OMP_NUM_THREADS=2.

6. Build twice, publish once

Key point

Reproducibility that nobody checks decays: one day a step starts writing a timestamp, and nobody notices until the day it matters. So the CI job that publishes the image also checks it.

On every pull request that touches the package it builds the image twice, the second time without any cache, and compares the two digests. On a release tag it does the same and pushes the first build to the GitHub Container Registry, ghcr.io, as ghcr.io/xmarva/perevod:<version>.

The image job in GitHub Actions. A pull request or a release tag starts it. It builds the image for x86-64 and ARM, then builds it again from scratch, and compares the two digests; if they differ, the job fails. On a release tag, the first build is pushed to ghcr.io/xmarva/perevod with the package's version as its tag The image job in GitHub Actions. A pull request or a release tag starts it. It builds the image for x86-64 and ARM, then builds it again from scratch, and compares the two digests; if they differ, the job fails. On a release tag, the first build is pushed to ghcr.io/xmarva/perevod with the package's version as its tag

Some details of the workflow below. SOURCE_DATE_EPOCH is the time of the commit, so a rebuild of the same commit next year uses the same time.

QEMU, an emulator, lets the x86 runner build the ARM image too, for people on Apple silicon and ARM servers. The builder is pinned, moby/buildkit:v0.34.1, because another version of BuildKit may pack the same files into different bytes.

Attestations are switched off: they record when and where a build ran, so no two builds have the same ones. The job logs in to the registry with the token GitHub gives every workflow run, and needs the packages: write permission for it [9].

write(".github/workflows/image.yml", """
name: image

# perevod's Docker image. On pull requests that touch the package: build it twice and check
# that the two builds are the same, bit for bit. On a release tag: the same, and publish it.
on:
  push:
    tags: ["perevod-v*"]
  pull_request:
    paths: ["perevod/**", ".github/workflows/image.yml"]

env:
  IMAGE: ghcr.io/xmarva/perevod
  PLATFORMS: linux/amd64,linux/arm64
  # OCI media types in both builds, and no attestations: those record when and where a build
  # ran, so no two builds have the same ones
  OUTPUT: oci-mediatypes=true,rewrite-timestamp=true
  PUSH: ${{ github.ref_type == 'tag' }}  # only a release tag publishes

jobs:
  image:
    runs-on: ubuntu-latest
    permissions:
      contents: read
      packages: write
    steps:
      - uses: actions/checkout@v7
      - uses: docker/setup-qemu-action@v4  # runs the arm64 build on an x86 machine
      - uses: docker/setup-buildx-action@v4
        with:
          driver-opts: image=moby/buildkit:v0.34.1  # a pinned builder writes the same bytes
      - name: Settings from the commit
        run: |
          echo "SOURCE_DATE_EPOCH=$(git log -1 --format=%ct)" >> "$GITHUB_ENV"
          init=perevod/src/perevod/__init__.py
          version=$(sed -n 's/^__version__ = "\\(.*\\)"$/\\1/p' "$init")
          echo "VERSION=$version" >> "$GITHUB_ENV"
          echo "REVISION=$GITHUB_SHA" >> "$GITHUB_ENV"
      - name: The tag matches the package's version
        if: github.ref_type == 'tag'
        run: test "$GITHUB_REF_NAME" = "perevod-v$VERSION"
      - uses: docker/login-action@v4
        if: github.ref_type == 'tag'
        with:
          registry: ghcr.io
          username: ${{ github.actor }}
          password: ${{ secrets.GITHUB_TOKEN }}
      - name: Build, and push on a release tag
        run: |
          docker buildx build perevod --platform "$PLATFORMS" \\
            --provenance=false --sbom=false \\
            --build-arg SOURCE_DATE_EPOCH --build-arg VERSION --build-arg REVISION \\
            --output "type=image,name=$IMAGE:$VERSION,push=$PUSH,$OUTPUT" \\
            --metadata-file first.json
      - name: Build again from scratch, and compare
        run: |
          docker buildx build perevod --platform "$PLATFORMS" \\
            --provenance=false --sbom=false --no-cache \\
            --build-arg SOURCE_DATE_EPOCH --build-arg VERSION --build-arg REVISION \\
            --output "type=oci,dest=second.tar,$OUTPUT" --metadata-file second.json
          first=$(jq -r '."containerimage.digest"' first.json)
          second=$(jq -r '."containerimage.digest"' second.json)
          echo "first build:  $first"
          echo "second build: $second"
          test "$first" = "$second"
""")

The README gets the commands to run the image, and the server’s, which 0.2.0’s README never mentioned:

replace("perevod/README.md", 'perevod "I think, therefore I am."',
        'perevod "Tom is afraid of the dark."')
replace("perevod/README.md", "The model was trained on short everyday sentences", """\
As a service, with an HTTP API at `/translate`, its documentation at `/docs`, and with
`--demo` a page at `/`:

```bash
pip install "perevod[demo] @ git+$REPO#subdirectory=perevod"
perevod-server --port 8000 --demo
```

Or from the image of a release, with the model inside, which needs no network to run:

```bash
docker run --rm -p 8000:8000 ghcr.io/xmarva/perevod:0.3.1
curl -X POST localhost:8000/translate -H "Content-Type: application/json" \\
    -d '{"text": "Hello!"}'
```

The model was trained on short everyday sentences""")

A release tag starts the workflow (its runs), and its log ends with two equal digests, or the job fails. The image it pushed is public, so the notebook can check the result itself. The registry is GitHub’s this time, with the same four requests:

VERSION = "0.3.1"
GHCR = "ghcr.io"
IMAGE = "xmarva/perevod"
ghcr_token = pull_token(GHCR, IMAGE)
response = get_manifest(GHCR, IMAGE, VERSION, ghcr_token)
print(f"{GHCR}/{IMAGE}:{VERSION}", response.headers["Docker-Content-Digest"])

manifests = {}
configs = {}
for entry in response.json()["manifests"]:
    arch = entry["platform"]["architecture"]
    manifests[arch] = get_manifest(GHCR, IMAGE, entry["digest"], ghcr_token).json()
    config_digest = manifests[arch]["config"]["digest"]
    configs[arch] = json.loads(get_blob(GHCR, IMAGE, config_digest, ghcr_token))
for name, value in configs["amd64"]["config"]["Labels"].items():
    print(name, "=", value)

ghcr.io/xmarva/perevod:0.3.1 sha256:8b46a084770cc8cd7aac71ca30ffd1c95807ad79af9200f5e7fb28464b27839b
org.opencontainers.image.revision = 1db63fa5a13dd09d00be9b7c6179850a1894500a
org.opencontainers.image.source = https://github.com/xmarva/transformers-from-scratch
org.opencontainers.image.version = 0.3.1

The labels say which commit the image was built from. And the two architectures can be compared layer by layer:

def short(command):
    command = " ".join(command.split())
    command = command.split(" /bin/sh -c ")[-1].replace(" # buildkit", "")
    return command[:56]


steps = []
for step in configs["amd64"]["history"]:
    if not step.get("empty_layer"):
        steps.append(step["created_by"])
rows = []
for k, step in enumerate(steps):
    x86 = manifests["amd64"]["layers"][k]
    arm = manifests["arm64"]["layers"][k]
    same = "yes" if x86["digest"] == arm["digest"] else "no"
    rows.append({"step": short(step), "MB (x86)": x86["size"] / 1e6, "same on ARM": same})
display(table(pd.DataFrame(rows), digits=1))
stepMB (x86)same on ARM
# debian.sh --arch 'amd64' out/ 'trixie' '@1792531975'29.8no
set -eux; apt-get update; apt-get install -y --no-instal1.6no
set -eux; savedAptMark="$(apt-mark showmanual)"; apt-get12.1no
set -eux; for src in idle3 pip3 pydoc3 python3 python3-c0.0no
useradd --uid 1000 --create-home perevod0.0no
WORKDIR /home/perevod0.0yes
COPY requirements.lock build.lock ./0.0yes
pip install --require-hashes --no-deps -r requirements.l285.2no
COPY pyproject.toml README.md ./src/0.0yes
COPY src ./src/src0.0yes
pip install --no-deps --no-build-isolation ./src && rm -0.1no
python -c "from perevod import Translator; Translator.fr35.9yes

Most layers differ, and should: Debian, Python and PyTorch are compiled for each processor. The interesting rows are the others. The copied files are the same bytes on both, and so is the model’s layer, which was made by running Python on each architecture, the ARM one inside an emulator, and keeping what the download wrote.

Two different processors, one emulated, produced the same bytes, because nothing in that layer depends on when or where it was made.

Two small layers of ours differ for a reason that took some digging: both carry times inherited from the base image, which was built separately, minutes apart, for each architecture.

useradd copies the new user’s start-up files from /etc/skel with their original times. And running pip makes Python import a few modules of its own standard library that the base image hadn’t cached yet, so Python writes their .pyc files into the layer that installs perevod.

Python’s import system writes timestamp headers no matter what SOURCE_DATE_EPOCH says, which only changes the default of py_compile and compileall, and the timestamps are those of the standard library’s files.

rewrite-timestamp only rewrites times later than SOURCE_DATE_EPOCH, and these are earlier. Inside one architecture they’re pinned along with the base, so two builds still agree.

There’s a stronger test than CI’s. CI builds twice on the same machine, minutes apart. The real promise of a reproducible build is that anyone, anywhere, later, gets the same bytes.

A day after the release, a rebuild of the same commit on a laptop, an ARM Mac, with the same command and the same pinned BuildKit, gave the published digest, on a fresh builder.

On a builder that had built earlier versions, one directory in one layer came out with the time of an old build, and the digest differed: rebuild on a fresh builder, as CI always does.

7. The image, without Docker

The notebook can’t run Docker, but it can do what Docker does when it pulls an image: download the layers, check each against its digest, and unpack them on top of each other into one folder. That folder is the image’s filesystem.

One detail matters while unpacking. A layer can’t change an earlier layer, so when a step deletes a file, its layer carries an empty marker named .wh.<name>, a whiteout, and the file must be removed from what’s been unpacked so far.

And a tar archive from the internet deserves suspicion: Python’s data filter refuses entries that would land outside the folder, including links to absolute paths, which would point into this machine’s own system instead of the image’s [10].

MACHINES = {"x86_64": ("amd64", "x86_64-linux-gnu", "ld-linux-x86-64.so.2"),
            "aarch64": ("arm64", "aarch64-linux-gnu", "ld-linux-aarch64.so.1")}
ARCH, TRIPLET, LOADER = MACHINES[os.uname().machine]
ROOT = WORK_DIR / "rootfs"
shutil.rmtree(ROOT, ignore_errors=True)
ROOT.mkdir()
refused = []


def unpack(member, path):
    name = Path(member.name)
    if name.name.startswith(".wh."):  # a whiteout: a later layer deleted this file
        gone = Path(path) / name.parent / name.name.removeprefix(".wh.")
        if gone.is_dir() and not gone.is_symlink():
            shutil.rmtree(gone)
        elif gone.exists() or gone.is_symlink():
            gone.unlink()
        return None
    try:
        return tarfile.data_filter(member, path)
    except tarfile.FilterError:
        refused.append(member.name)
        return None


start = time.time()
downloaded = 0
for layer in manifests[ARCH]["layers"]:
    data = get_blob(GHCR, IMAGE, layer["digest"], ghcr_token)  # checked against its digest
    downloaded += len(data)
    with tarfile.open(fileobj=io.BytesIO(data), mode="r:gz") as tar:
        tar.extractall(ROOT, filter=unpack)
seconds = time.time() - start
print(f"{downloaded / 1e6:.0f} MB downloaded, checked and unpacked in {seconds:.0f} s")
print(len(refused), "links refused, for example", refused[:3])

365 MB downloaded, checked and unpacked in 37 s
185 links refused, for example ['etc/alternatives/awk', 'etc/alternatives/awk.1.gz', 'etc/alternatives/builtins.7.gz']

The refused entries are Debian’s links such as /etc/alternatives/awk, which point to absolute paths. They only make sense inside the container, and the program doesn’t need them. Look around:

print(sorted(path.name for path in ROOT.iterdir()))
site = ROOT / "usr/local/lib/python3.12/site-packages"
print(len(list(site.glob("*.dist-info"))), "installed packages")
models = ROOT / "home/perevod/models"
for path in sorted(models.rglob("*")):
    if path.is_file() and "snapshots" in path.parts:
        print(path.relative_to(models), f"{path.stat().st_size:,} bytes")
leftovers = []
if (ROOT / "home/perevod/.cache").exists():
    leftovers = list((ROOT / "home/perevod/.cache").rglob("*"))
print("anything in ~/.cache:", leftovers or "no")

['bin', 'boot', 'dev', 'etc', 'home', 'lib', 'lib64', 'media', 'mnt', 'opt', 'proc', 'root', 'run', 'sbin', 'srv', 'sys', 'tmp', 'usr', 'var']
41 installed packages
models--xmarva--perevod-en-ru/snapshots/4abc0d63e862ca99d2bd99cd1007e4a9fe47c358/config.json 235 bytes
models--xmarva--perevod-en-ru/snapshots/4abc0d63e862ca99d2bd99cd1007e4a9fe47c358/generation_config.json 58 bytes
models--xmarva--perevod-en-ru/snapshots/4abc0d63e862ca99d2bd99cd1007e4a9fe47c358/model.safetensors 38,453,000 bytes
models--xmarva--perevod-en-ru/snapshots/4abc0d63e862ca99d2bd99cd1007e4a9fe47c358/tokenizer.json 1,230,871 bytes
anything in ~/.cache: no

A Linux system, Python’s packages, and the model at the commit perevod pins. Is it the model that’s on the Hub? The Hub stores the sha256 of every large file, so ask it, and hash the file from the image:

from huggingface_hub import HfApi

MODEL = "xmarva/perevod-en-ru"
REVISION = "4abc0d63e862ca99d2bd99cd1007e4a9fe47c358"
info = HfApi().model_info(MODEL, revision=REVISION, files_metadata=True)
for file in info.siblings:
    if file.rfilename == "model.safetensors":
        hub_sha = file.lfs.sha256
snapshot = models / "models--xmarva--perevod-en-ru" / "snapshots" / REVISION
weights = snapshot / "model.safetensors"
image_sha = hashlib.sha256(weights.read_bytes()).hexdigest()
print("on the Hub:   ", hub_sha)
print("in the image: ", image_sha)
print("the same:", hub_sha == image_sha)

on the Hub:    0b085d1e77815f47b17cbd5b400ec735edff9a222fb9e18c32c00101168e021d
in the image:  0b085d1e77815f47b17cbd5b400ec735edff9a222fb9e18c32c00101168e021d
the same: True

Now the cause that the table in section 4 missed. perevod 0.3.0 passed CI’s check, two identical builds, and its model layer still held a file that had no business there. Open that layer of the 0.3.0 image:

old_index = get_manifest(GHCR, IMAGE, "0.3.0", ghcr_token).json()
for entry in old_index["manifests"]:
    if entry["platform"]["architecture"] == ARCH:
        old_manifest = get_manifest(GHCR, IMAGE, entry["digest"], ghcr_token).json()
old_model_layer = get_blob(GHCR, IMAGE, old_manifest["layers"][-1]["digest"], ghcr_token)
with tarfile.open(fileobj=io.BytesIO(old_model_layer), mode="r:gz") as tar:
    for member in tar.getmembers():
        if member.isfile() and "/hub/" not in member.name:
            print(member.name, member.size, "bytes")
    found = tar.extractfile("home/perevod/.cache/huggingface/.agent_harnesses.json")
    agents = json.loads(found.read())
names = list(agents["harnesses"])
print(len(names), "AI coding agents, starting with", names[:4])

home/perevod/.cache/huggingface/.agent_harnesses.json 5778 bytes
25 AI coding agents, starting with ['antigravity', 'augment-cli', 'cline', 'cowork']

Next to the model sat a list of AI coding agents. huggingface_hub 2.x labels its requests with the agent that made them, if any, and to recognise agents it downloads the list from the Hub, “at most once a day”, and caches it in ~/.cache/huggingface.

The list lives on the Hub’s server and changes when the Hub adds an agent. So the two builds of the release agreed, but a rebuild of the same commit after the next change to that list would not: the model layer would hold a different file and get a different digest.

Nothing pins a file the build downloads from a server, except the build itself. perevod 0.3.1 keeps the model in a folder of its own and deletes everything else under ~/.cache in the same step, which is the line you saw in the Dockerfile.

Key point

The lesson in it is general: a step that runs a program records everything that program writes, not only what you asked for. Look inside a layer at least once.

Now run the image. A container is, at heart, a process that sees the image’s folder as its whole filesystem.

Docker arranges that with the kernel’s namespaces, which a notebook can’t use, but there’s a simpler trick that’s enough here: run the image’s own Python through the image’s own dynamic loader, the program that starts every Linux binary and loads its libraries.

Pointed at the image’s libraries, it runs the image’s Python with the image’s C library and the image’s packages, whatever this machine has installed. The environment comes from the image’s config, with its paths moved into the folder:

config = configs[ARCH]["config"]
env = {}
for item in config["Env"]:
    name, value = item.split("=", 1)
    if value.startswith("/home/"):
        value = str(ROOT) + value  # a path inside the image
    env[name] = value
env["PATH"] = os.environ["PATH"]
env["HOME"] = str(ROOT / "home/perevod")
env["PYTHONHOME"] = str(ROOT / "usr/local")
env["HF_HUB_DISABLE_PROGRESS_BARS"] = "1"
libraries = [ROOT / "usr/local/lib", ROOT / "usr/lib" / TRIPLET]
PYTHON = [str(ROOT / "usr/lib" / TRIPLET / LOADER), "--library-path",
          ":".join(str(path) for path in libraries), str(ROOT / "usr/local/bin/python3.12")]
for name in ["HF_HUB_OFFLINE", "HF_HUB_CACHE", "PYTHON_VERSION"]:
    print(name, "=", env[name])

probe = """
import sys, torch, perevod
from perevod import Translator
translator = Translator.from_pretrained(device="cpu")
print("Python", sys.version.split()[0], "| torch", torch.__version__,
      "| perevod", perevod.__version__)
print("model:", translator.source)
"""
result = subprocess.run([*PYTHON, "-c", probe], env=env, capture_output=True, text=True)
print(result.stdout.strip() or result.stderr[-2000:])

HF_HUB_OFFLINE = 1
HF_HUB_CACHE = /kaggle/working/tfs/rootfs/home/perevod/models
PYTHON_VERSION = 3.12.16
Python 3.12.16 | torch 2.14.1+cpu | perevod 0.3.1
model: xmarva/perevod-en-ru@4abc0d63e862ca99d2bd99cd1007e4a9fe47c358

Python 3.12 from the image, its PyTorch, its perevod, and the model loaded from inside it with the Hub switched off. This notebook’s own Python is a different version, with a different PyTorch, which makes a good experiment: install the package from the source the notebook has been editing, translate the same sentences both ways, and compare.

PIP_INSTALL = [sys.executable, "-m", "pip", "install", "-q"]
subprocess.run([*PIP_INSTALL, "./perevod[test,serve]"], check=True)
importlib.invalidate_caches()

import perevod
from perevod import Translator, decoding

DATA_REVISION = "540be5b93bf77cd0846b02d55f2558a5596887a9"
test_path = hf_hub_download("xmarva/tfs-tatoeba-en-ru", "data/test.parquet",
                            repo_type="dataset", revision=DATA_REVISION)
sentences = pd.read_parquet(test_path)["en"].sample(200, random_state=SEED).tolist()

translator = Translator.from_pretrained(device="cpu")
here = translator.translate(sentences)
print("this notebook: Python", sys.version.split()[0], "| torch", torch.__version__,
      "| perevod", perevod.__version__)

job = """
import json, sys
from perevod import Translator
sentences = json.load(sys.stdin)
translations = Translator.from_pretrained(device="cpu").translate(sentences)
print(json.dumps(translations, ensure_ascii=False))
"""
result = subprocess.run([*PYTHON, "-c", job], env=env, input=json.dumps(sentences),
                        capture_output=True, text=True, check=True)
there = json.loads(result.stdout)
same = 0
for a, b in zip(here, there, strict=True):
    if a == b:
        same += 1
print(f"identical translations: {same} of {len(sentences)}")

this notebook: Python 3.13.15 | torch 2.11.0+cpu | perevod 0.3.1
identical translations: 200 of 200

A different Python and a different PyTorch, and here every translation came out the same. That isn’t guaranteed.

PyTorch’s kernels change between versions, and so does the order of their arithmetic; when two hypotheses in a beam are nearly tied, a difference in the last digit of a score picks a different one. The image is how you stop wondering whether that will happen to you.

Finally the image’s own command, the server, started the same way, and asked the question that 0.3 added:

(On Kaggle and Colab nothing else is listening on port 8009; on your machine, pick a free port.)

start_server = "from perevod.server import main; main(['--port', '8009'])"
server = subprocess.Popen([*PYTHON, "-c", start_server], env=env,
                          stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
for _ in range(120):  # up to a minute for the server to start
    try:
        health = requests.get("http://127.0.0.1:8009/health", timeout=1).json()
        break
    except requests.ConnectionError:
        time.sleep(0.5)
question = {"text": "Tom is afraid of the dark."}
answer = requests.post("http://127.0.0.1:8009/translate", json=question)
server.terminate()
server.wait()
print(health)
print(answer.json())

{'status': 'ok', 'version': '0.3.1', 'model': 'xmarva/perevod-en-ru@4abc0d63e862ca99d2bd99cd1007e4a9fe47c358'}
{'translation': 'Том боится темноты.'}

That’s the image working, minus what Docker adds on top: its own process tree, network and hostname (namespaces), limits on CPU and memory (cgroups), and a root filesystem the process can’t see past.

Our trick shares all of those with the notebook, which is fine for a check and no way to run a service. The unpacked image takes more than a gigabyte, so it goes:

shutil.rmtree(ROOT)

8. The same model in a browser

An image needs a machine to run on, and someone to pay for the machine. For a demo there’s a cheaper place to run a model this small: the visitor’s own device.

The obvious host for a free demo used to be a Hugging Face Space with Gradio on a free CPU. That changed in July 2026: now “Gradio and Docker Spaces run on compute and require a paid plan to create”, while “Static Spaces are free for everyone”[11].

A static Space serves files, HTML and JavaScript, and nothing stops one of those files from being a model. ONNX Runtime Web runs ONNX models in the browser with WebAssembly, and perevod’s weights take 38 MB. So the page downloads the model once, and every translation runs on the visitor’s processor: no server, and nothing typed leaves the device.

The in-browser demo. A static Hugging Face Space serves the page, the tokenizer as JSON and two ONNX files: the encoder and one decoder step. In the visitor's browser, a JavaScript tokenizer turns the text into ids, ONNX Runtime Web runs the encoder once and the decoder step once per beam search step, and the beam search in JavaScript picks the translation The in-browser demo. A static Hugging Face Space serves the page, the tokenizer as JSON and two ONNX files: the encoder and one decoder step. In the visitor's browser, a JavaScript tokenizer turns the text into ids, ONNX Runtime Web runs the encoder once and the decoder step once per beam search step, and the beam search in JavaScript picks the translation

ONNX is a file format for a computation graph: the operations of a network and its weights, readable by many runtimes. The model goes in as two graphs. The encoder runs once per sentence.

The decoder step takes the hypotheses so far and returns the log-probabilities of the next token for each, so the beam search of lesson 3 stays outside the graph, written in JavaScript, where it’s a loop of forty lines.

The step has no KV cache, so each step recomputes the whole prefix; for short sentences that costs less than it sounds, and keeps the graph simple.

The export is a script in the course repository, next to the page’s files:

write("space/export_onnx.py", '''
# Export perevod's model to the two ONNX graphs of the in-browser demo:
#   encoder.onnx: source ids -> memory
#   decoder.onnx: hypotheses, memory, source ids -> log-probabilities of the next token
# and the tokenizer, which the page loads as it is.
# Usage: python export_onnx.py --out model
import argparse
import logging
from pathlib import Path

import torch
from torch.export import Dim

from perevod import Translator, decoding


class Encoder(torch.nn.Module):
    def __init__(self, model):
        super().__init__()
        self.model = model

    def forward(self, src):
        return self.model.encode(src)


class DecoderStep(torch.nn.Module):
    def __init__(self, model):
        super().__init__()
        self.model = model

    def forward(self, tgt, memory, src):
        return self.model.decode(tgt, memory, src)[:, -1].log_softmax(-1)


def main():
    parser = argparse.ArgumentParser(description="Export perevod's model to ONNX")
    parser.add_argument("--out", default="model")
    out = Path(parser.parse_args().out)
    out.mkdir(parents=True, exist_ok=True)
    logging.getLogger("torch.onnx").setLevel(logging.ERROR)  # the exporter is chatty
    translator = Translator.from_pretrained(device="cpu")
    model = translator.model
    _, bos, _ = decoding.special_ids(translator.tokenizer)
    ids = decoding.source_ids(translator.tokenizer, "Tom is afraid of the dark.", 256)
    src = torch.tensor([ids, ids])  # two rows, three target tokens: any sizes above one do
    tgt = torch.tensor([[bos, bos, bos], [bos, bos, bos]])
    batch = Dim("batch")
    source = Dim("source", max=256)  # the model's positions
    target = Dim("target", max=256)
    # export leaves a module in the mode it found it in, and a new module starts in training
    # mode: without eval() the wrappers would switch the model's dropout back on afterwards
    encoder = Encoder(model).eval()
    step = DecoderStep(model).eval()
    with torch.no_grad():
        memory = model.encode(src)
        torch.onnx.export(encoder, (src,), out / "encoder.onnx", input_names=["src"],
                          output_names=["memory"], external_data=False, verbose=False,
                          dynamic_shapes={"src": {0: batch, 1: source}})
        torch.onnx.export(step, (tgt, memory, src), out / "decoder.onnx",
                          input_names=["tgt", "memory", "src"], output_names=["logp"],
                          external_data=False, verbose=False,
                          dynamic_shapes={"tgt": {0: batch, 1: target},
                                          "memory": {0: batch, 1: source},
                                          "src": {0: batch, 1: source}})
    translator.tokenizer.save(str(out / "tokenizer.json"))  # the page needs it too
    for path in sorted(out.glob("*.onnx")):
        print(f"{path.name}: {path.stat().st_size / 1e6:.1f} MB")


if __name__ == "__main__":
    main()
''')

ONNX = ["onnx==1.19.1", "onnxscript==0.7.2", "onnxruntime==1.31.1"]
subprocess.run([*PIP_INSTALL, *ONNX], check=True)
result = subprocess.run([sys.executable, "space/export_onnx.py", "--out", "space/model"],
                        capture_output=True, text=True, check=True)
print(result.stdout.strip())

decoder.onnx: 29.6 MB
encoder.onnx: 26.3 MB

The two graphs share the embeddings in PyTorch, so together they’re larger than the weights: each file carries its own copy. Dim tells the exporter which sizes may change between calls: the number of hypotheses, and the lengths of the source and of the translation so far.

Before anything goes near a browser, check that the graphs compute what PyTorch computes. ONNX Runtime has a Python package with the same engine:

import onnxruntime as ort

encoder = ort.InferenceSession("space/model/encoder.onnx")
decoder = ort.InferenceSession("space/model/decoder.onnx")
pad_id, bos_id, eos_id = decoding.special_ids(translator.tokenizer)

src = torch.tensor([decoding.source_ids(translator.tokenizer, sentences[0], 256)])
tgt = torch.tensor([[bos_id]])
with torch.no_grad():
    memory = translator.model.encode(src)
    expected = translator.model.decode(tgt, memory, src)[:, -1].log_softmax(-1).numpy()
onnx_memory = encoder.run(None, {"src": src.numpy()})[0]
got = decoder.run(None, {"tgt": tgt.numpy(), "memory": onnx_memory, "src": src.numpy()})[0]
print("largest difference in log-probabilities:", float(np.abs(got - expected).max()))
print("the same first token:", int(got.argmax()) == int(expected.argmax()))

largest difference in log-probabilities: 7.62939453125e-06
the same first token: True

The differences are in the fifth or sixth decimal place: the same arithmetic, done in a different order. What matters is the translations. Here is lesson 3’s beam search once more, written against the two graphs the way the page’s JavaScript does it, step for step: every hypothesis is fed again whole, and the four best continuations survive.

def length_penalty(length, alpha=0.6):
    return ((5 + length) / 6) ** alpha


def onnx_translate(sentence, beam_size=4, max_tokens=128):
    ids = decoding.source_ids(translator.tokenizer, sentence, 256)
    src = np.array([ids], dtype=np.int64)
    memory = encoder.run(None, {"src": src})[0]
    beams = [([bos_id], 0.0)]
    finished = []
    for _ in range(max_tokens):
        n = len(beams)
        tgt = np.array([tokens for tokens, score in beams], dtype=np.int64)
        feeds = {"tgt": tgt, "memory": np.repeat(memory, n, axis=0),
                 "src": np.repeat(src, n, axis=0)}
        logp = decoder.run(None, feeds)[0]
        candidates = []
        for (tokens, score), row in zip(beams, logp):
            for token in np.argsort(-row, kind="stable")[:beam_size]:
                candidates.append((tokens + [int(token)], score + float(row[token])))
        candidates.sort(key=lambda candidate: candidate[1], reverse=True)
        beams = []
        for tokens, score in candidates:
            if tokens[-1] == eos_id:
                finished.append((tokens, score))
            else:
                beams.append((tokens, score))
            if len(beams) == beam_size:
                break
        if len(finished) >= beam_size or not beams:
            break
    pool = finished or beams
    best, best_score = pool[0][0], -math.inf
    for tokens, score in pool:
        normalized = score / length_penalty(len(tokens) - 1)
        if normalized > best_score:
            best, best_score = tokens, normalized
    return translator.tokenizer.decode([token for token in best[1:] if token != eos_id])


start = time.time()
through_onnx = []
for sentence in sentences:
    through_onnx.append(onnx_translate(sentence))
same = 0
for a, b in zip(here, through_onnx, strict=True):
    if a == b:
        same += 1
seconds = time.time() - start
print(f"identical to the package: {same} of {len(sentences)}, in {seconds:.0f} s")
for sentence, translation in list(zip(sentences, through_onnx))[:3]:
    print(sentence, "->", translation)

identical to the package: 200 of 200, in 35 s
The novel has been translated into German. -> Роман был переведён на немецкий.
I wonder why you never told me about it. -> Интересно, почему ты никогда мне об этом не рассказывал.
Arguing is no way to go about it. -> И спорить об этом нет.

Word for word. The page itself is three small files in the course repository’s space/ folder, and only one of them is about the page.

translate.js is the function above in JavaScript. tokenizer.js rewrites the tokenizer of lesson 1, whose tokenizer.json the page loads as it is, and it had one trap worth knowing.

Trap

The tokenizer splits text into words with the regular expression \w+, and in the Rust library behind Hugging Face’s tokenizers, \w matches letters of every alphabet. In JavaScript it matches only [A-Za-z0-9_], even with the Unicode flag, so a word like “naïve” or “Москва” would come apart.

The JavaScript version spells the classes out, [\p{Alphabetic}\p{M}\p{Nd}\p{Pc}\p{Join_Control}], and a script in the same folder checks it against the original on every English sentence of the test set, id for id.

The rest is a static Space: a README.md whose header says sdk: static, uploaded with the two graphs:

python space/export_onnx.py --out space/model
hf upload xmarva/perevod space . --repo-type space

Here it is, xmarva/perevod:

The perevod page on Hugging Face Spaces: an English text box with 'My brother lives in a small town by the sea.', the Russian translation 'Мой брат живёт в маленьком городке у моря.' below it, the time the translation took on this device, a note that the model runs on the visitor's device and nothing typed leaves it, and four example sentences

The first visit downloads the two graphs once, and the browser keeps them. After that the page works without a network, and the translation takes about as long as in this notebook, on whatever device it opened on. The model can’t do more there than here, of course; it’s the same model, with the same limits its card lists.

Two ways to ship one model, then, for two jobs. The image is for running it as a service: the same bytes on every server, checkable by anyone, with nothing left to chance at startup.

The page is for showing it to people: free, private, and limited by the visitor’s device. Both are built from the same weights at the same commit, and both were checked against the package, translation for translation.

9. Release 0.3.1

Before a release, the fast tests, which now include the health check that names the model:

!cd perevod && {sys.executable} -m pytest -q -p no:cacheprovider -m "not slow" --color=no

........................                                                 [100%]
24 passed, 2 deselected in 6.97s

The release is a commit and a tag, perevod-v0.3.1, and the image is whatever CI builds from them. That gives the notebook one last check: are the files it wrote and edited the same as the release’s?

Only the files that go into the image count, the build context and the workflow. The lock files are compared without their first lines, where uv records the command, which here includes --exclude-newer:

RELEASE = "perevod-v0.3.1"
subprocess.run(["git", "fetch", "-q", "--depth", "1", "origin", "tag", RELEASE], check=True)
listed = subprocess.run(["git", "ls-tree", "-r", "--name-only", RELEASE, "perevod"],
                        capture_output=True, text=True, check=True).stdout.split()
IMAGE_FILES = {"Dockerfile", ".dockerignore", "requirements.lock", "build.lock",
               "pyproject.toml", "README.md"}
context = [".github/workflows/image.yml"]
for name in listed:
    if name.startswith("perevod/src/") or name.removeprefix("perevod/") in IMAGE_FILES:
        context.append(name)


def without_command(text):
    lines = []
    for line in text.splitlines():
        if not line.startswith("#    uv pip compile"):
            lines.append(line)
    return lines


different = []
for name in context:
    released = subprocess.run(["git", "show", f"{RELEASE}:{name}"], capture_output=True,
                              text=True, check=True).stdout
    if without_command(Path(name).read_text()) != without_command(released):
        different.append(name)
print(len(context), "files compared, different:", different or "none")

15 files compared, different: none

The notebook rebuilt the release’s inputs file for file, and CI turned them into ghcr.io/xmarva/perevod:0.3.1, the image section 7 took apart. On your machine:

docker run --rm -p 8000:8000 ghcr.io/xmarva/perevod:0.3.1
curl localhost:8000/health

0.3.1 is a patch release in lesson 7’s terms: the package’s code is the same as in 0.3.0, and only the way the image is built changed. Pinned by its digest, the image will translate the same way in a year, on any machine that runs Docker, with or without a network. That was the point of the lesson.

References

  1. NumPy 2.0.0 release notes, and the error in practice: PyTorch issue #129310. ↩
  2. Python Packaging User Guide: install_requires vs requirements files. ↩
  3. Stufft, D. (2013). setup.py vs requirements.txt. Abstract and concrete dependencies. ↩
  4. pip: secure installs and hash-checking mode; uv: uv pip compile and using uv with PyTorch. ↩
  5. docker-library FAQ and Docker’s build best practices, on tags and digests. ↩
  6. reproducible-builds.org: the definition, and SOURCE_DATE_EPOCH. ↩
  7. BuildKit: reproducible builds; Docker: image exporter options. ↩
  8. PEP 552: Deterministic pycs; py_compile. ↩
  9. GitHub: publishing Docker images and the Container registry. ↩
  10. Python: tarfile extraction filters; the dynamic loader: ld.so(8). ↩
  11. Hugging Face: Spaces overview and static Spaces. ↩
  12. Open Container Initiative: image specification (index, manifest, config, layers, descriptors) and distribution specification.
  13. huggingface_hub environment variables (HF_HUB_CACHE, HF_HUB_OFFLINE).
  14. PyTorch: exporting to ONNX; ONNX Runtime Web.
]]> Eva Koroleva Is it rational to be good? Game theory and the ethics of cooperation 2025-12-11T12:00:00+00:00 2026-11-02T00:00:00+00:00 https://blog.korolevalab.com/blog/2026/game-theory-ethics

In short

Glaucon’s challenge in the Republic, that people are just only for fear of being wronged and that a ring of invisibility would make anyone unjust, has the structure of the prisoner’s dilemma: whatever the other player does, defecting pays more. Game theory answers him in part.

For people who expect to meet again, who are watched by others and who cannot be sure that their rings will keep working, justice is very often the rational choice, and the models even say which kind: nice, firm, forgiving within limits, and clear.

They do not answer the question Glaucon actually asked, whether justice is good in itself for someone whose ring never fails.

In the second book of Plato’s Republic, Glaucon offers Socrates a challenge that moral philosophy has been answering ever since. People, he says, praise justice not for itself but for its consequences.

Left to ourselves we would all prefer to wrong others and get away with it; we settle for justice because we fear being wronged more than we enjoy wronging.

Justice, on this account, is “a mean or compromise, between the best of all, which is to do injustice and not be punished, and the worst of all, which is to suffer injustice without the power of retaliation” (Republic 359a)[1].

To make the point he tells the story of Gyges, a shepherd who finds a ring that makes him invisible, and uses it to seduce the queen, kill the king and take the throne.

Give such a ring to a just man and to an unjust one, Glaucon says, and they will end up doing the same things, since “no man can be imagined to be of such an iron nature that he would stand fast in justice”, and “wherever any one thinks that he can safely be unjust, there he is unjust” (360b–c).

Glaucon’s argument has a structure, and the structure has a name. Twenty-three centuries after the Republic, the mathematicians of game theory found it in a two-by-two table and called it the prisoner’s dilemma.

The table turned the old question into a precise one: under which conditions is it rational to be just, and which kind of justice is rational?

This essay follows the answers that game theory gives, from Hobbes to evolutionary simulations, and marks the point where they stop. The models are small enough to compute exactly, and every picture below is computed from them.

1. The prisoner’s dilemma

Two people can each do something for the other at a cost to themselves: keep a promise, share credit, return a favour. If both do, both are better off.

If one does and the other doesn’t, the one who didn’t is best off of all and the one who did is worst off. If neither does, both lose what they could have had together.

Game theory calls this structure the prisoner’s dilemma, after a story Albert Tucker told in 1950 to illustrate a game that Merrill Flood and Melvin Dresher had invented a few months earlier at the RAND Corporation[2].

The prisoner's dilemma with Axelrod's payoffs. Both cooperate: 3 points each. You cooperate and they defect: 0 for you, 5 for them. You defect and they cooperate: 5 for you, 0 for them. Both defect: 1 each. Whatever the other does, defecting gives you more, and two players who reason this way get 1 each The prisoner's dilemma with Axelrod's payoffs. Both cooperate: 3 points each. You cooperate and they defect: 0 for you, 5 for them. You defect and they cooperate: 5 for you, 0 for them. Both defect: 1 each. Whatever the other does, defecting gives you more, and two players who reason this way get 1 each

Only the order of the four outcomes matters. The temptation to exploit beats the reward for mutual cooperation, which beats the punishment for mutual defection, which beats the sucker’s payoff:

\[T > R > P > S, \qquad 2R > T + S.\]

The second condition says that taking turns at exploiting each other is worse than cooperating, so that cooperation, not alternation, is what the two players lose. The numbers in the table, \(T = 5\), \(R = 3\), \(P = 1\), \(S = 0\), are the ones Robert Axelrod used, and the rest of this essay uses them too.

The reasoning that follows is short and, within its assumptions, impeccable. Whatever the other player does, you do better by defecting: \(5 > 3\) if they cooperate, \(1 > 0\) if they don’t. Defection is a dominant strategy, and two players who see this end up with 1 point each, while 3 each was on the table.

The outcome is a Nash equilibrium, since neither can gain by changing alone, and it is worse for both than another outcome they could have reached. Individual rationality, applied correctly by both, produces a result that both regard as a failure.

defecting pays whatever the other doesboth players defect1 point each instead of 3

This is Glaucon’s compromise in miniature. Each would most like to wrong the other with impunity; each fears most to be wronged; and justice, mutual cooperation, is an agreement they would both welcome and neither has a reason to keep.

2. Hobbes: changing the game

Thomas Hobbes described the same structure as the condition of humanity without government.

In the state of nature, where nobody can trust anybody to hold back, everyone has reason to strike first, and the result is a war of all against all, in which there is “no place for Industry; because the fruit thereof is uncertain”, and the life of man is “solitary, poore, nasty, brutish, and short” (Leviathan, ch. 13)[3].

Hobbes also gave Glaucon’s position its sharpest voice. “The Foole hath sayd in his heart, there is no such thing as Justice”, and he argues that “to make, or not make; keep, or not keep Covenants, was not against Reason, when it conduced to ones benefit” (ch. 15). The Foole is the player who has noticed that defection dominates.

Hobbes’s reply is not that the Foole makes a logical error in the single case. It is that the Foole is not in a single case.

A man who breaks his covenant, and declares that he may with reason do so, “cannot be received into any Society, that unite themselves for Peace and defence, but by the errour of them that receive him”, and “if he be left, or cast out of Society, he perisheth” (ch. 15).

The Foole’s defection is a single move, but it tells everyone what kind of player he is, and the rest of the game is played against that knowledge.

Hobbes’s remedy for the war of all against all, though, is not an argument but an institution:

Covenants, without the Sword, are but Words, and of no strength to secure a man at all.

Thomas Hobbes, Leviathan, ch. 17

The sovereign does not persuade the players to cooperate; he changes their payoffs, so that defection no longer pays. In the language of game theory, Hobbes’s solution is to make the game stop being a prisoner’s dilemma.

That works, and modern states, courts and contracts are its descendants. But it leaves Glaucon’s question untouched, because it explains why people keep promises when someone watches them, which Glaucon never doubted. His question was about the ring.

3. Can cooperation in a single encounter be rational?

Some philosophers have argued that the dominance reasoning itself is wrong, or at least incomplete, even when the two players meet once and no sovereign watches.

The first argument is Kantian. Kant’s categorical imperative reads, in Thomas Kingsmill Abbott’s translation, “Act only on that maxim whereby thou canst at the same time will that it should become a universal law” (Groundwork, Ak 4:421)[4].

One of Kant’s own examples is close to the dilemma: a man borrows money with a promise he knows he will not keep.

If everyone acted on that maxim, “the promise itself would become impossible, as well as the end that one might have in view in it, since no one would consider that anything was promised to him” (Ak 4:422).

The defector’s maxim, accept the cooperation of others and give none back, has the same shape. It works only as an exception, and made universal it destroys the cooperation it lives on.

In 1983 the cognitive scientist Douglas Hofstadter turned a related thought into an argument about the game itself. His superrational thinkers, “by recursive definition, include in their calculations the fact that they are in a group of superrational thinkers”[5].

Two such players know that they will reach the same conclusion, so the only outcomes really open to them are the two on the diagonal, and of those mutual cooperation is better.

David Lewis showed that this is a special case of a deeper split in decision theory: “Prisoners’ Dilemma is a Newcomb Problem—or rather, two Newcomb Problems side by side, one per prisoner”[6].

In Newcomb’s problem, a very reliable predictor has put a million dollars into an opaque box if it predicted that you would take only that box, and nothing if it predicted that you would also take a transparent box with a thousand dollars in it.

Against a player who reasons as you do, the other player’s choice is like the predictor’s: your choice does not cause it, but it is very good evidence about it.

Evidential decision theory, which chooses the act that is the best news about the outcome, recommends taking one box, and cooperating.

Causal decision theory, which chooses the act with the best causal consequences, recommends taking both boxes, and defecting, because nothing you do now changes what is in the box or what the other player does. Philosophers have not agreed which theory is right, and the disagreement is now half a century old.

David Gauthier took a different route[7]. Rather than ask which act is rational, he asked which disposition it is rational to have. A straightforward maximizer always chooses the act with the best expected outcome; a constrained maximizer cooperates with others who are disposed to cooperate, and defects otherwise.

If dispositions are translucent, if others can tell, at least sometimes, what kind of person you are, constrained maximizers are offered cooperation and straightforward maximizers are not, and the constrained ones do better. The rational choice of a disposition can then require acts that, taken one at a time, are not rational.

Critics have pointed out that the argument works only to the extent that dispositions really are visible, which brings Gauthier close to the next idea: that what makes cooperation rational is that the encounter is not, in fact, single.

4. The shadow of the future

David Hume saw the problem in an agricultural setting (Treatise 3.2.5.8)[8]:

Your corn is ripe to-day; mine will be so to-morrow. ‘Tis profitable for us both, that I shou’d labour with you to-day, and that you shou’d aid me to-morrow. I have no kindness for you, and know you have as little for me. I will not, therefore, take any pains upon your account; and should I labour with you upon my own account, in expectation of a return, I know I shou’d be disappointed, and that I shou’d in vain depend upon your gratitude. Here then I leave you to labour alone: You treat me in the same manner. The seasons change; and both of us lose our harvests for want of mutual confidence and security.

The farmers in Hume’s example are not wicked; they are prudent, and prudence ruins them both. What saves them is not a sovereign but the expectation of further dealings.

“Hence I learn to do a service to another, without bearing him any real kindness; because I forsee, that he will return my service, in expectation of another of the same kind” (3.2.5.9).

The promise, on Hume’s account, is a convention that makes this expectation public: whoever promises “subjects himself to the penalty of never being trusted again in case of failure” (3.2.5.10).

Game theory makes this precise. Suppose that after each round there is another with probability \(w\), so that a relationship lasts \(1/(1 - w)\) rounds on average. Now a strategy can depend on what the other player did.

The simplest such strategy, tit for tat, cooperates first and then copies the other player’s previous move. Against it, a defection buys \(T\) today and costs a punishment tomorrow, and it pays to cooperate if and only if

\[w \;\ge\; \max\!\left(\frac{T - R}{T - P},\; \frac{T - R}{R - S}\right) = \frac{2}{3}.\]

The first term guards against always defecting, the second against defecting and making up in turns[9]. Axelrod called \(w\) the shadow of the future[10].

Points per round against tit for tat as a function of w, the probability of another round. Cooperating earns 3. Always defecting earns 5 - 4w and wins below w=1/2; defecting and making up in turns earns 5/(1 + w) and wins below w=2/3. Above 2/3 cooperating is the best reply Points per round against tit for tat as a function of w, the probability of another round. Cooperating earns 3. Always defecting earns 5 - 4w and wins below w=1/2; defecting and making up in turns earns 5/(1 + w) and wins below w=2/3. Above 2/3 cooperating is the best reply

The most striking illustration comes from the First World War. On quiet stretches of the Western Front, British and German battalions faced each other for weeks at a time and, without negotiating anything, stopped trying to kill each other.

A British staff officer was “astonished to observe German soldiers walking about within rifle range behind their own line. Our men appeared to take no notice … These people evidently did not know there was a war on. Both sides apparently believed in the policy of ‘live and let live’“[11, 10].

When a salvo landed by mistake, a German climbed onto his parapet and shouted: “We are very sorry about that; we hope no one was hurt. It is not our fault, it is that damned Prussian artillery.”

The truce ended when headquarters ordered raids, whose results it could check: “There was no effective way to pretend that a raid had been undertaken when it had not.”

Two philosophical lessons follow, and the second is less comfortable than the first. The first is that the rationality of cooperation is a fact about the structure of a relationship as much as about the people in it. The same person has more reason to keep a promise to a neighbour than to a stranger at an airport.

The second is that the repeated game does not select cooperation; it only permits it. The theorems that show cooperation can be an equilibrium of a repeated game, known collectively as the folk theorem, also show that almost any pattern of behaviour can be one, including mutual defection forever and many unequal arrangements in between.

Key point

Rationality alone underdetermines what rational people will do together. Something else, a convention, a norm, a shared history, has to pick one equilibrium among many, and that something is where much of morality lives.

5. Which character is rational?

If many strategies can sustain cooperation, which ones do well? Axelrod asked game theorists, economists, psychologists and mathematicians to submit programs and played them against each other in a round robin[12].

The winner was “the simplest of all strategies submitted”, tit for tat, sent by the psychologist Anatol Rapoport[9]. In a second round, 62 entrants who knew this tried to beat it, and it won again[13].

A rerun of that second tournament from Axelrod’s original code with 25,000 random seeds has since found that tit for tat wins about 63% of them[14], so the result is a tendency, not a law. Here is a small tournament of the same kind, with twelve classic strategies:

Matches won against mean points per round for twelve strategies. Generous ZD scores highest, 2.50, winning 0 matches; the nice strategies take the top places without winning any. Always defect wins 9 of 11 matches and scores 2.05; Joss is last with 1.87 Matches won against mean points per round for twelve strategies. Generous ZD scores highest, 2.50, winning 0 matches; the nice strategies take the top places without winning any. Always defect wins 9 of 11 matches and scores 2.05; Joss is last with 1.87

The six strategies that never defect first take the top six places, while always defect wins nine of its eleven matches and finishes ninth.

Between them, those six win almost no individual matches. The Extortioner, one of the zero-determinant strategies that William Press and Freeman Dyson found in 2012, is built never to earn less than its partner[15].

It loses no match and finishes eleventh. Its generous twin, built never to earn more than its partner, wins no match and finishes first[16].

Axelrod’s explanation is that a nice strategy can never get more than its partner in a single game, since it only defects in answer to a defection; it wins the tournament by doing reasonably well with everyone.

His advice “Don’t be envious” is a point about which comparison is rational: your score is measured against others in your position, not against the person across the table.

Axelrod summed up the winning character in four words, “nice, retaliatory, forgiving, and clear”[10]. Each word names something a moral philosopher would recognise: benevolence, a readiness to resent injury, a willingness to let it go, and the transparency that makes a person predictable.

What game theory adds is that the four belong together. Niceness without retaliation is exploited, retaliation without forgiveness turns every accident into a feud, and both without clarity are wasted, because the other side cannot learn from them.

The feud is easy to see once mistakes are allowed, as they always are in life:

Sixteen rounds between two players of the same strategy, with one mistake in round 4: player one means to cooperate and defects. Two tit for tat players then punish each other in turns for the rest of the game. Two generous tit for tat players stop the echo when one of them forgives, here after 2 rounds. Two win-stay, lose-shift players defect together once and are back to cooperating in round 6 Sixteen rounds between two players of the same strategy, with one mistake in round 4: player one means to cooperate and defects. Two tit for tat players then punish each other in turns for the rest of the game. Two generous tit for tat players stop the echo when one of them forgives, here after 2 rounds. Two win-stay, lose-shift players defect together once and are back to cooperating in round 6

Between two tit for tat players, one misread gesture starts an exchange of punishments that never ends, each side certain that it is only responding. A generous version of tit for tat, which forgives a defection with some probability \(q\), stops the echo after a few rounds[17].

How large should \(q\) be? Forgive too rarely, and mistakes turn into feuds; forgive too often, and a player who alternates defection and apology earns more than honest cooperators. The largest forgiveness that no exploiter can profit from is

\[q^* = \min\!\left(1 - \frac{T - R}{R - S},\; \frac{R - P}{T - P}\right) = \frac{1}{3}\]

with these payoffs[18, 17]:

Left: residents who forgive a defection with probability q earn 3; an invader who defects every other round earns more once q passes 1/3, one who always defects once q passes 1/2. Right: with 5 percent mistakes the residents earn more the more they forgive Left: residents who forgive a defection with probability q earn 3; an invader who defects every other round earns more once q passes 1/3, one who always defects once q passes 1/2. Right: with 5 percent mistakes the residents earn more the more they forgive

It is hard not to read this as a formal version of Aristotle’s doctrine of the mean.

“Virtue, then, is a state of character concerned with choice, lying in a mean, i.e. the mean relative to us”; to feel anger “at the right times, with reference to the right objects, towards the right people, with the right motive, and in the right way, is what is both intermediate and best” (Nicomachean Ethics II.6)[19].

Forgiveness here has an excess and a deficiency, and its mean is not a midpoint but depends on the situation, as Aristotle said it should: as the temptation \(T\) grows relative to the reward \(R\), the right amount of forgiveness falls.

6. Evolution and the origin of morality

If cooperative dispositions pay, natural selection and cultural imitation will spread them, and biologists have taken this as an explanation of where morality comes from.

Robert Trivers opened the paper that introduced reciprocal altruism with an admission: “Models that attempt to explain altruistic behavior in terms of natural selection are models designed to take the altruism out of altruism”[20].

Here is what imitation does to a population of a hundred players who copy strategies that earn more, with relationships of twenty rounds on average and one move in fifty going wrong:

Shares of strategies in a population of a hundred over time. Without generous tit for tat the population cycles: defectors are replaced by tit for tat, which drifts into unconditional cooperation, which defectors take over again; over twenty runs defectors hold the majority 49% of the time. With generous tit for tat it mostly holds the population, and defectors hold the majority 18% of the time Shares of strategies in a population of a hundred over time. Without generous tit for tat the population cycles: defectors are replaced by tit for tat, which drifts into unconditional cooperation, which defectors take over again; over twenty runs defectors hold the majority 49% of the time. With generous tit for tat it mostly holds the population, and defectors hold the majority 18% of the time

With only always defect, tit for tat and unconditional cooperation available, the population never settles. Tit for tat drives out defectors; unconditional cooperators then drift in, since among reciprocators they do just as well; once there are enough of them, defectors return.

Lorens Imhof, Drew Fudenberg and Martin Nowak described this cycle in finite populations[21]. Over twenty runs of the simulation, defectors held the majority 49% of the time.

Unconditional kindness, on this picture, is not a stable state of a society but a phase of it, and what protects a cooperative society is a willingness to punish that erodes exactly when it seems least necessary.

When generous tit for tat is available, it takes over and holds the population: defectors held the majority 18% of the time, and the average payoff rose from 1.94 to 2.58 points per round.

The evolutionary story raises an old worry in a new form. If our moral dispositions exist because they paid, are they really moral?

One answer is that the explanation of a disposition’s origin is not an account of its content: a person who keeps a promise because they value honesty is honest, whatever selection did to their ancestors.

Another answer comes from the way cooperation among strangers actually works. Richard Alexander argued that “moral systems are systems of indirect reciprocity” (p. 77)[22]: I help you, and someone else helps me, because they saw what I did.

Reputation then does the work that repetition does between two people, and it requires judging others, which brings normative questions back in. Martin Nowak and Karl Sigmund showed that cooperation can be sustained by a public “image score” that rises when you help and falls when you refuse[23].

But a player who refuses to help a known cheat then loses reputation for enforcing the norm, so a rule called standing, which looks at whom you refused, works much better[24]:

How an observer updates a donor's reputation. Image scoring looks only at the act: helping makes the donor good, refusing makes the donor bad. Standing also looks at whom the donor refused: refusing someone with a bad reputation keeps the donor good. The two rules differ only in that one case How an observer updates a donor's reputation. Image scoring looks only at the act: helping makes the donor good, refusing makes the donor bad. Standing also looks at whom the donor refused: refusing someone with a bad reputation keeps the donor good. The two rules differ only in that one case

When Hisashi Ohtsuki and Yoh Iwasa checked all 4096 rules of this kind, only eight kept cooperation stable, and all eight share two features: helping a good person is good and refusing one is bad; refusing a bad person is good, “because it works as sanction”[25].

The evolution of cooperation selects not only for helping but for a particular moral judgement: that punishing the guilty is just and punishing the punisher is not.

Indirect reciprocity also gives an answer to Hume’s sensible knave, who reasons: “That honesty is the best policy, may be a good general rule; but is liable to many exceptions: And he, it may, perhaps, be thought, conducts himself with most wisdom, who observes the general rule, and takes advantage of all the exceptions” (Enquiry 9.22)[26].

Hume admitted that if a man thinks this reasoning requires an answer, “it will be a little difficult to find any, which will to him appear satisfactory and convincing” (9.23).

The answer he did give was about the knave’s inner life: knaves “are, in the end, the greatest dupes, and have sacrificed the invaluable enjoyment of a character, with themselves at least, for the acquisition of worthless toys and gewgaws” (9.25).

The game-theoretic answer is about his outer life. In a world of indirect reciprocity, the knave’s exceptions are exactly what others observe and remember, and only a knave who could be sure of never being seen would profit from them. That knave is Gyges again, and rings like his are rare.

7. The stag hunt and the social contract

Not every problem of cooperation is a prisoner’s dilemma. Rousseau told a story in which nobody gains by betraying anyone: “Was a deer to be taken? Every one saw that to succeed he must faithfully stand to his post; but suppose a hare to have slipped by within reach of any one of them, it is not to be doubted but he pursued it without scruple”[27].

Brian Skyrms argued that this game deserves at least as much attention as the prisoner’s dilemma: “The Stag Hunt is a story that became a game. The game is a prototype of the social contract”[28].

Expected catch of hunting the stag (4 times the share p of others who hunt it) and of hunting hare (3). Hunting the stag pays when more than three quarters of the others hunt it; below that the share of stag hunters falls, above it it rises Expected catch of hunting the stag (4 times the share p of others who hunt it) and of hunting hare (3). Hunting the stag pays when more than three quarters of the others hunt it; below that the share of stag hunters falls, above it it rises

In the stag hunt, mutual cooperation is an equilibrium: nobody wants to leave a successful hunt. But mutual distrust is an equilibrium too, and it is the safer one, what John Harsanyi and Reinhard Selten called risk-dominant[29].

With a stag worth 4 to each hunter, a hare worth 3 and nothing for hunting the stag alone, hunting the stag pays only if more than three quarters of the others will hunt it too.

The problem is not temptation but trust, and trust here is self-fulfilling: the same people, with the same values, end up in either equilibrium depending on what they expect of each other.

This is the structure of what Hume called convention, and his example is the cleanest: “Two men, who pull the oars of a boat, do it by an agreement or convention, tho’ they have never given promises to each other” (Treatise 3.2.2.10)[8].

David Lewis built his analysis of convention on such cases, Hume’s rowers among them: a regularity is a convention if everyone conforms to it, everyone expects everyone else to conform, and everyone prefers to conform on condition that the others do (p. 42)[30].

On this view, a good deal of what we call morality, the rules of fair division, of turn-taking, of truthfulness in ordinary talk, is less a constraint on self-interest than a way of coordinating on one of many equilibria.

Skyrms noted the link with the shadow of the future: “If the probability of repetition is high enough, the stag hunting equilibrium becomes risk dominant.”

8. What game theory cannot settle

Every model in this essay takes the payoffs as given and asks only how to play. Amartya Sen argued that this is where the theory is weakest: a person who acts only on their own preferences, with no room for commitment to others or to principles, is not a model of rationality but of a kind of impoverishment.

In his words, “The purely economic man is indeed close to being a social moron” (p. 336)[31]. Experiments agree that real players are not like this.

In the ultimatum game, where one player proposes how to split a sum and the other can accept or let both get nothing, the narrowly rational offer is the smallest possible one, and typical offers are 40 to 50 percent[32].

Across fifteen small-scale societies, mean offers ranged from 26 to 58 percent, and “the higher the degree of market integration and the higher the payoffs to cooperation, the greater the level of cooperation in experimental games”[33].

Even the trench truce has been questioned as an example: Andrew Gelman argued that not shooting simply lowered each side’s own risk, so the soldiers may not have faced a prisoner’s dilemma at all[34].

This sets the limit of the whole approach. Game theory can show that, given what people want, certain dispositions are rational: niceness, retaliation, a measured forgiveness, attention to reputation, a readiness to trust when trust is likely to be returned.

Key point

Game theory cannot show what people should want. The payoffs are inputs, and the question of what is worth wanting, which is most of ethics, comes before the first row of the table.

What, then, of Glaucon? Game theory answers him in part. For people who expect to meet again, who are watched by others, and who cannot be sure that their rings will keep working, justice is very often the rational choice, and the models even say which kind of justice: nice, firm, forgiving within limits, and clear.

They do not answer the question Glaucon actually asked, which was whether justice is good in itself, for someone whose ring never fails. That question the models leave where Plato found it.

References

  1. Plato. Republic, book II, 358e–360d. Translated by B. Jowett. Project Gutenberg #1497. ↩
  2. Flood, M. M. (1952). Some Experimental Games. RAND Research Memorandum RM-789-1. ↩
  3. Hobbes, T. (1651). Leviathan, chapters 13, 15 and 17 (pp. 62, 72–73 and 85 of the first edition). Project Gutenberg #3207. ↩
  4. Kant, I. (1785). Groundwork of the Metaphysics of Morals, Ak 4:421–422. Translated by T. K. Abbott as Fundamental Principles of the Metaphysic of Morals. Project Gutenberg #5682. ↩
  5. Hofstadter, D. R. (1983). Dilemmas for superrational thinkers, leading up to a luring lottery. Scientific American, 248(6), 14–28. Reprinted in Metamagical Themas (1985), ch. 30. ↩
  6. Lewis, D. (1979). Prisoners’ dilemma is a Newcomb problem. Philosophy & Public Affairs, 8(3), 235–240. ↩
  7. Gauthier, D. (1986). Morals by Agreement. Clarendon Press. Chapter VI. ↩
  8. Hume, D. (1739–40). A Treatise of Human Nature, book 3, part 2, sections 2 and 5. Hume Texts Online. ↩
  9. Axelrod, R., & Hamilton, W. D. (1981). The evolution of cooperation. Science, 211(4489), 1390–1396. ↩
  10. Axelrod, R. (1984). The Evolution of Cooperation. Basic Books. Chapter 4, The live-and-let-live system in trench warfare in World War I, pp. 73–87. ↩
  11. Ashworth, T. (1980). Trench Warfare 1914–1918: The Live and Let Live System. Macmillan. ↩
  12. Axelrod, R. (1980a). Effective choice in the prisoner’s dilemma. Journal of Conflict Resolution, 24(1), 3–25. ↩
  13. Axelrod, R. (1980b). More effective choice in the prisoner’s dilemma. Journal of Conflict Resolution, 24(3), 379–403. ↩
  14. Knight, V., Campbell, O., Harper, M., Gaffney, T. J., & Glynatsi, N. E. (2025). Reviving, reproducing, and revisiting Axelrod’s second tournament. arXiv:2510.15438. ↩
  15. Press, W. H., & Dyson, F. J. (2012). Iterated Prisoner’s Dilemma contains strategies that dominate any evolutionary opponent. PNAS, 109(26), 10409–10413. ↩
  16. Stewart, A. J., & Plotkin, J. B. (2013). From extortion to generosity, evolution in the Iterated Prisoner’s Dilemma. PNAS, 110(38), 15348–15353. ↩
  17. Nowak, M., & Sigmund, K. (1992). Tit for tat in heterogeneous populations. Nature, 355, 250–253. ↩
  18. Molander, P. (1985). The optimal level of generosity in a selfish, uncertain environment. Journal of Conflict Resolution, 29(4), 611–618. ↩
  19. Aristotle. Nicomachean Ethics, book II. Translated by W. D. Ross. ↩
  20. Trivers, R. L. (1971). The evolution of reciprocal altruism. Quarterly Review of Biology, 46(1), 35–57. ↩
  21. Imhof, L. A., Fudenberg, D., & Nowak, M. A. (2005). Evolutionary cycles of cooperation and defection. PNAS, 102(31), 10797–10800. ↩
  22. Alexander, R. D. (1987). The Biology of Moral Systems. Aldine de Gruyter; reissued by Routledge, 2017. ↩
  23. Nowak, M. A., & Sigmund, K. (1998). Evolution of indirect reciprocity by image scoring. Nature, 393, 573–577. ↩
  24. Leimar, O., & Hammerstein, P. (2001). Evolution of cooperation through indirect reciprocity. Proceedings of the Royal Society B, 268, 745–753. ↩
  25. Ohtsuki, H., & Iwasa, Y. (2004). How should we define goodness? Reputation dynamics in indirect reciprocity. Journal of Theoretical Biology, 231(1), 107–120. ↩
  26. Hume, D. (1751). An Enquiry concerning the Principles of Morals, section 9, part 2. Hume Texts Online. ↩
  27. Rousseau, J.-J. (1755). A Discourse upon the Origin and the Foundation of the Inequality among Mankind, part II. ↩
  28. Skyrms, B. (2001). The stag hunt. Proceedings and Addresses of the American Philosophical Association, 75(2), 31–41. ↩
  29. Harsanyi, J. C., & Selten, R. (1988). A General Theory of Equilibrium Selection in Games. MIT Press. ↩
  30. Lewis, D. (1969). Convention: A Philosophical Study. Harvard University Press. The introduction and chapters I–II in facsimile, Econ Journal Watch, 16(2), 2019. ↩
  31. Sen, A. K. (1977). Rational fools: a critique of the behavioral foundations of economic theory. Philosophy & Public Affairs, 6(4), 317–344. ↩
  32. Güth, W., & Kocher, M. G. (2014). More than thirty years of ultimatum bargaining experiments. Journal of Economic Behavior & Organization, 108, 396–409. ↩
  33. Henrich, J., Boyd, R., Bowles, S., Camerer, C., Fehr, E., Gintis, H., & McElreath, R. (2001). In search of homo economicus: behavioral experiments in 15 small-scale societies. American Economic Review, 91(2), 73–78. ↩
  34. Gelman, A. (2008). Methodology as ideology: some comments on Robert Axelrod’s The Evolution of Cooperation. QA Rivista dell’Associazione Rossi-Doria, 2/2008, 167–176. ↩

Further reading

The models, the simulations and the figures: articles/game_theory.

]]> Eva Koroleva