SpeedrunningPLMs: the benchmark data

Train a protein language model in ten minutes on one GPU, and know exactly how good it is. This is the fixed data of SpeedrunningPLMs, a masked language model speedrun on proteins: 1,754,034 training proteins from the AlphaFold Database as one token stream, and held-out splits whose masks are drawn once, so every run, every seed and every reference model is scored on the same hidden residues. No held-out protein has a training relative at 30 percent sequence identity over 80 percent coverage, by sequence (MMseqs2) or by structure (Foldseek), so a lower loss means a model that generalizes across protein families, not one that remembers them.

The benchmark id eba8e0f953ffead6 is the SHA-256 of manifest.json; every result carries it, and two results compare only when it agrees.

What a run does with it

Step	Rule
Corruption	15 percent of the standard residues are replaced by MASK, nothing else; CLS, EOS and padding are never masked. Training uses the same corruption.
Loss	Cross entropy over the 20 standard residues at the masked positions, in nats per masked residue. A model with a larger vocabulary (ESM2 has 33 tokens, ESMC 64) is scored on its 20 residue columns, so every model faces the same 20-way problem.
Clock	Ten minutes of applied training steps on one GPU. Compile warmup and in-run validation are outside the clock and capped.
Selection	The weights with the lowest loss on val_select (one protein per validation family, at most 1,024 tokens) are scored once on the full validation and test splits at full length (at most 2,048 tokens). The headline is the test loss.
Structure probe	A linear probe from the hidden states to the 20 Foldseek 3Di states, trained on the fixed 20,000-protein subset and scored on validation and test. Reported beside the loss, never used for selection.

Splits

Split	Proteins	Families	Residues	Scored (masked) residues	Max tokens
train	1,754,034	333,143	389,084,153	drawn each step	free
test	25,778	4,807	6,098,489	915,208	2,048
val_select	4,751	4,751	1,066,372	156,231	1,024
validation	25,459	4,751	5,919,918	889,145	2,048

Families are disjoint across the three splits, inherited from the Alignment Atlas family split of the AlphaFold Database (Synthyra/alignment_atlas_families_afdb_v6 at 83737fa33e72). val_select is a subset of validation: each family's longest representative, truncated, for the in-run selection; the full splits are scored once at the end. Lengths run from 50 to 2,046 residues, and every result reports the loss by length bin (50 to 99, 100 to 249, 250 to 499, 500 to 999, 1,000 and more).

Where models stand

Test loss on this data, every model scored by the same evaluator on the same masked residues. The reference models trained for days on UniRef, which shares sequences with these proteins, so their rows are a comparator rather than a fair baseline; the ten-minute record and the full table live in the code repository.

Model	Trained on	Test loss
unigram residue frequencies	the floor	2.8933
ESM2 8M	UniRef50, days	2.6253
ESM2 35M	UniRef50, days	2.4992
ESM2 150M	UniRef50, days	2.3303
ESM2 650M	UniRef50, days	2.0710
ESM2 3B	UniRef50, days	1.8705
ESMC 300M	UniRef and more, days	2.1269
ESMC 600M	UniRef and more, days	2.0510
ESMC 6B	UniRef and more, days	1.6891

Loading

pip install git+https://github.com/Synthyra/SpeedrunningPLMs

from pathlib import Path

from huggingface_hub import snapshot_download
from speedrunning_plms import benchmark

data_dir = Path(snapshot_download("Synthyra/SpeedrunningPLMs", repo_type="dataset"))
stream = benchmark.load_train(data_dir)                 # memory-mapped tokens, offsets, 3Di states, probe subset
protein = stream.protein(0)                            # (length + 2,) uint8: CLS, residues, EOS
test = benchmark.load_held_out(data_dir, "test")       # tensors sorted by descending length, digests verified
metrics = benchmark.evaluate_split(model, test, device)  # any model exposing logits and its 20 residue columns

The files need no library to read: the streams are flat NumPy arrays (np.memmap("train_tokens.u8", dtype=np.uint8)) and the held-out files are torch.load dictionaries. Token ids are ESM2's: CLS 0, PAD 1, EOS 2, MASK 32, and the residues A 5, C 23, D 13, E 9, F 18, G 6, H 21, I 12, K 15, L 4, M 20, N 17, P 14, Q 16, R 10, S 8, T 11, V 7, W 22, Y 19. The 3Di classes index Foldseek's alphabet ACDEFGHIKLMNPQRSTVWY in order, with 255 at positions that are not residues.

Files

File	Size	What it holds
train_tokens.u8	374 MB	one uint8 ESM2 token stream of every training protein, CLS, residues, EOS, in a seed-0 shuffled order
train_offsets.i64	13 MB	int64: where each protein starts in the stream, plus the stream's end
train_three_di.u8	374 MB	uint8: the Foldseek 3Di class of every stream position, 255 at CLS and EOS
train_families.json	23 MB	the family of each training protein, in stream order
probe_subset.i64	156 KB	int64: the stream indices of the structure probe's 20,000 training proteins, one per family
val_select.pt	14 MB	the selection split: input_ids, masked, three_di, lengths, family, families
validation.pt	145 MB	the full validation split, the same tensors at full length
test.pt	143 MB	the full test split, the same tensors at full length
manifest.json	small	the source, the objective, the seeds, every count and every file's SHA-256

Every file's SHA-256 is in manifest.json and the loaders verify it, so a run cannot read altered data unnoticed.

Source and terms

AlphaFold Database sequences (CC BY 4.0) and their Foldseek 3Di states, grouped into families by sequence and structure and split by family in Synthyra's Alignment Atlas. This release carries the tokens, the 3Di states and the family index of each protein, no coordinates. The dataset is CC BY 4.0; cite the code repository and the AlphaFold Database.