A 4.94M-parameter masked-language-model protein encoder trained on eukaryotic sequences. I built it to study what integer quantisation does to a protein language model — and, along the way, to have a small fast encoder that is actually usable.
It matches ESM2-8M on contextual structure while running roughly twice as fast at 66% of the parameters.
ESM/BERT-style: RoPE, pre-norm, SDPA (FlashAttention during inference), d_model 320, 8 heads, 6 layers, d_ff 640, predicting the 20 standard amino acids at masked positions. Trained under bf16 autocast, mmseqs at 50% identity using a nextflow pipeline.
python bin/downloads.py eukaryota
python bin/downloads.py proteingym
./nextflow run main.nfmain.nf PREP -> TRAIN -> {EVAL_PGYM, QUANTISE, LAYERWISE, CONTEXT, QAT} -> REPORT + FIGURES
nextflow.config all parameters + SLURM/Singularity resources
Dockerfile
bin/
core.py model, data pipeline, metrics, quantisation helpers
train.py training and QAT fine-tuning
downloads.py fetch training FASTA, ProteinGym, ubiquitin probe set
quantise.py PTQ ladder, per-layer sensitivity, mixed precision, precision baseline
proteingym.py zero-shot variant-effect scoring vs ESM2, with quantised variants
context_ablation.py delta_context: NLL(shuffled context) - NLL(intact)
layerwise.py tuned-lens per-layer NLL
rank_check.py effective rank of the representation
homology.py training-set homology, for the memorisation objection
bench.py latency, memory and throughput across precisions
packed.py real n-bit packed weights (int/e2m1/nf4)
report.py leaderboard + champion selection
plots/ every figure in make_figures.sh
flashee/ pip-installable inference wrapper (torch only, no training stack)
hf/README.md model card for the Hugging Face Hub
Numbers below are the an undertrained model. Higher scale training happening rn. More diagrams coming.
delta_context = NLL with the context shuffled minus NLL with it intact, on identical probe positions.
ESM2-8M 0.3040
ours 0.3003 98.8%
untrained 0.0000
ProteinGym (198 assays): 0.1445 mean Spearman vs ESM2's 0.2118. The gap comes from the proteins the model never learned.
The gap closes with evolutionary proximity: 81% of ESM2 on eukaryotes, 79% on human, 47% on prokaryotes, 31% on viruses.
The undertrained model has effective rank of circa 11 compared to ESM2's 15. Will be seen if that is still the case with a fully trained models.
RTX A4500, batch 64, sequence length 512, bf16:
params weights ms/seq peak MiB
ours 4.94M 9.46 MiB 0.319 219.6
ESM2-8M 7.51M 14.37 MiB 0.665 444.1
2.06x 2.02x
At batch 1 the dtype makes no difference at all (fp32 3.73 ms vs bf16 3.75) — a 5M-parameter encoder is kernel-launch bound, not arithmetic bound.
Training and evaluation both run under autocast, so the model is bf16 native. Every compression ratio here is quoted against bf16, not fp32, which halves the headline numbers relative to the usual convention.
The ladder: int8 is essentially free, int4 costs little, int3 is a soft edge where GPTQ still holds but round-to-nearest starts to cliff, and int2 collapses to random baseline. QAT experiments on the way.
ff 46.0% embedding 19.0% W_v 15.2%
W_o 7.9% W_k 7.1% fc 2.5% W_q 2.3%
All six query projections together cost 1.12 top-1 points — 12% of the parameters, essentially free, which is consistent with softmax being robust to perturbed logits. By layer the damage is U-shaped: layers 0 and 5 carry 40% of it and the four middle layers 39%, independently reproducing what the tuned-lens probe says about where depth is doing work.
Weight-only packing reduces the storage needed, but does not boost speed or memory (!). Against bf16 at batch 64, int4 packed weights are 3.77× smaller but 3.6× slower and use 1.8× more peak memory, because dequantisation to floating point dominates and activations hold the memory. At batch 32 the weights are 9% of peak. For a model this size, quantisation is a great for compression, but does not boost inference speed, in line with literature.
Optimised using Weights and Biases using a sweep (./sweep.yaml)
A previous iteration of this project, a fungal-only encoder (~9.5M parameters), was helpful at learning things that informed design decisions with this project.
- BLOSUM62 correlation does not necessarely correlate with fitness. The model reached 0.36 Spearman against a 0.13 frequency-null, but a higher BLOSUM score over 0.3 did not result in super accurate ProteinGym scores.
- Fitness is gated by calibration, even within a single protein. On ubiquitin the fungal model beat ESM2 on 2 of 3 assays (Mavor 2016, Roscoe 2013), and the number of wins across seeds rose with how far each seed drove wild-type NLL down (3/3 at NLL 1.3, 2/3 at 2.2, none at 2.8). This is what motivated the wt-NLL stratification used throughout.
- Distillation is expensive (!) and often inefficient. ESM2 distillation lifted token metrics but added no fitness signal and eroded the ubiquitin advantage; the teacher's representation space also differs enough to needed an adapter (my one was of dubious quality so I was not certain if ESM2 representations matched the one of my model).
- Tiny models are not great at contact preditctions lol. While it achieved 2–3× above the random baseline, it still stayed below 0.10.
- Clustering (and data selected for training the model) is extremely important!. 50% identity dedup beat looser settings on ProteinGym; looser clustering raises MLM top-1 through train/test homolog leakage while lowering downstream quality.
The fungal weights and their inference wrapper (fungalplm) still live on HuggingFace if anyones interested.
Maciej Szczesny, as part of Stracquadanio Lab, 2026.