Built flashEE from scratch in PyTorch, a 4.94M-parameter protein language model trained on eukaryotic UniProt sequences, covering the data pipeline, architecture, training and evaluation. To the best of my knowledge, this is the smallest protein language model out there.
It matches ESM2-8M on eukaryotic ProteinGym fitness assays while running 2.2× faster and using 2.7× less memory per sequence, on two-thirds of the parameters. Implemented every major 4-bit quantisation format to reduce the model size down to 2.5MiB, and checked how well it preserves the model.
Generously supported by CCPBioSim and EPSRC
Weights: szchesny/flashee
pip install flasheefrom flashee import FlashEE
plm = FlashEE.load("szchesny/flashee")
emb = plm.embed(["MQIFVKTLTGKTITLEVEPSDTIENVK"]) # [N, 320] mean-pooled
res = plm.embed(seqs, per_residue=True) # list of [Li, 320]python bin/downloads.py eukaryota
python bin/downloads.py proteingym
./nextflow run main.nfProtein language models have inherited the scaling playbook of natural language processing: more data, more parameters, more compute. The result is a family of models that are expensive to run and, for many practical tasks, far larger than the task requires. In this project, I built a small protein language model, Flash Eukaryotic Encoder (flashEE), with 4.94M parameters, trained only on eukaryotic UniProt sequences, and used it to test which natural language processing (NLP) efficiency techniques genuinely transfer to small PLMs.
flashEE matches ESM2-8M on eukaryotic ProteinGym fitness assays while running up to 2.2x faster and using 2.7x less memory per sequence, on 66% of the parameters (4.94M against ESM2-8M's 7.51M), while performing worse on prokaryotic and viral assays, which were not included in the training dataset.
Additionally, quantisation, often used to improve efficiency, typically did not yield faster speeds or reduced memory usage at this scale. Still, four-bit weight quantisation is near-lossless in quality, reaches full parity after QAT, and shrinks the model by roughly 8x. On CPU, the speed gain of dynamic int8 falls within measurement noise at every batch size tested. In larger batch sizes, relevant in bigger pipelines, memory is mainly consumed by activations. 8-bit activation quantisation is lossless, whereas 4-bit retains only 37.5% of the original quality, making W4A8 (4-bit weights with 8-bit activations) the minimal deployable configuration.
I ran all experiments on a cluster with an NVIDIA A4500 GPU and managed the training process and all subsequent analyses through a Nextflow pipeline with 3 different seeds.
I set out to determine if a smaller model, trained on a more specialised dataset, could match the performance of similar-sized models. To do this, I enlarged the training set from just a few fungal proteomes to include most eukaryotic protein sequences available in UniProt.
Data selection was tricky: I initially filtered the UniProt database for non-broken strands under 512 (the maximum length that fit in my model's attention matrix during training), but that produced poor results. Many motifs were only present in longer proteins. To counter this, I used a technique used during ESM-C training: randomly trimming the protein so it fits the attention matrix. This drastically improved my results and immediately yielded higher scores. Sequences were then packed into batches by sorted length, so that similar-length proteins share a batch and padding waste is minimised.
The model is a 6-layer, 8-head encoder-only masked language model with 320 dimensions, using pre-norm, rotary positional embeddings (RoPE), and PyTorch's native scaled dot-product attention, trained using BERT-style masking. Major differences from ESM2: my model uses a custom embedder and a 29-token vocabulary instead of 33, reducing the number of distinct "ideas" the embedding must represent and improving the interpretability of learned features. Additionally, I halved the feed-forward network (FFN) dimensions from 4x to 2x. During a hyperparameter sweep in Weights & Biases, models with a smaller FFN reached a lower validation loss at the same step count, going against the conventional 4x setting in "Attention Is All You Need." This reduction decreased the parameter count from 7.5M to under 5M. The sweep also showed that the best training results occurred without dropout, so I removed it.
I implemented a strict normalisation protocol by adding norm layers before major calculations to enhance training stability and prevent memory spikes during large computations. Additionally, I minimised external packages beyond PyTorch. I didn't import Biopython, as I only used it to load FASTA sequences (easily handled with a simple script) and the BLOSUM62 matrix. I also didn't import the FlashAttention package due to compatibility issues with other software and hardware. Instead, I used torch's native Scaled Dot-Product Attention (SDPA), which runs FlashAttention when possible.
The model trained for 200k steps under mixed precision, at an effective batch size of 512 (batch 32 with gradient accumulation of 16).
Protein language model embeddings are difficult to interpret. To check how well my model predicts reality, I used a variety of tests.
The first metric I used was perplexity, the exponentiated cross-entropy, which is by far the easiest to calculate. flashEE reaches 10.3 with a masked top-1 accuracy of 29.8%, against a marginal-frequency baseline of 9.5%. On the BLOSUM62 matrix, which identifies which amino acids the model confuses, just as evolution does, flashEE scores ρ = 0.496 against a baseline of 0.158.
One of the most useful and fastest-to-calculate metrics I used was
Using this measure, we can determine how much of the final prediction is context-dependent rather than random; with the baseline at 0, flashEE scores 0.399 against ESM2-8M's 0.304.
ProteinGym is the standard benchmark for zero-shot variant effect prediction, organised by organism and functional assay. flashEE scores 0.1718 ± 0.0013 against ESM2-8M's 0.2111 across all 201 assays, but it is dragged down by assays for organisms it hasn't seen, while performing on par with ESM2-8M on eukaryotic and human tasks.
| By organism | n | flashEE | ESM2-8M | Relative |
|---|---|---|---|---|
| Eukaryote | 39 | 0.2269±0.0072 | 0.2230 | +1.7%±3.3 |
| Human | 87 | 0.2286±0.0026 | 0.2384 | −4.1%±1.1 |
| Prokaryote | 49 | 0.1011±0.0047 | 0.2059 | −50.9%±2.3 |
| Virus | 26 | 0.0326±0.0061 | 0.1116 | −70.8%±5.5 |
Performance was uneven across functions: flashEE beats ESM2-8M by 37% on organismal fitness but falls 18% short on stability assays.
| Function | n | flashEE | ESM2-8M | Relative |
|---|---|---|---|---|
| Organismal fitness | 33 | 0.2165±0.0088 | 0.1579 | +37.1%±5.6 |
| Activity | 27 | 0.2093±0.0065 | 0.2185 | −4.2%±3.0 |
| Expression | 16 | 0.2573±0.0048 | 0.2705 | −4.9%±1.8 |
| Binding | 10 | 0.2858±0.0166 | 0.3010 | −5.1%±5.5 |
| Stability | 40 | 0.2241±0.0078 | 0.2748 | −18.4%±2.8 |
On the other hand, it does better than ESM2-8M on proteins resembling those it was trained on, which is what you would expect from a specialist. Inspired by Hou et al. 2026, I sorted the assays by the model's wild-type NLL on the target protein to assess how well it learned each function.
| Protein NLL | n | flashEE | ESM2-8M | Relative |
|---|---|---|---|---|
| < 1.5 (learned well) | 22–26 | 0.3869±0.0093 | 0.3219±0.0081 | +20.2% |
| 1.5 – 2.3 (partially learned) | 27–35 | 0.3636±0.0113 | 0.2627±0.0051 | +38.5% |
| > 2.3 (not learned) | 144–149 | 0.0974±0.0046 | 0.1825±0.0016 | −46.6% |
A linear probe on each layer showed that almost all of the model's skill appears at the very end, with only layers 5 and 6 contributing significantly to the final representation. On top of that, effective rank analysis showed that flashEE occupies a much smaller subspace than ESM2-8M, with only ~13 dimensions pooled and 25 per residue, against ESM2-8M's roughly 63 per residue.
Contact prediction (P@L) was conducted, and both flashEE and the ESM model performed poorly at this scale (~10% correct, barely above baseline). The same holds for reconstructing eukaryote phylogenetic relationships, with both models unable to predict families at this scale.
flashEE's peak memory is 39.8 MiB, plus 2.33 MiB per sequence, compared with 43.3 MiB + 6.26 MiB for ESM2-8M. The fixed portion mainly comprises the weights, while the per-sequence component relates to activations.
The first major optimisation involved modifying the rotary positional embeddings (RoPE), which encode each residue's position. By storing the rotation tables at half width, broadcasting them, and integrating the rotation into a single kernel, we reduced activation memory by 7% and runtime by 6.6% over bf16, from 114.4 to 109.4 MiB at B=32. In addition, running the model with weights in the bf16 format resulted in even faster inference.
| Variant (B=32) | ms/seq | vs fp32 | Peak MiB | MiB/seq |
|---|---|---|---|---|
| fp32 | 1.126±0.009 | 1.00× | 179.0 | 4.69 |
| bf16 | 0.325±0.002 | 3.47×±0.01 | 114.4 | 2.33 |
| bf16 + fused RoPE | 0.305±0.002 | 3.70×±0.01 | 109.4 | 2.17 |
| ESM2-8M, fp32 | 1.925 | — | 459.0 | 12.51 |
| ESM2-8M, bf16 | 0.700 | — | 243.7 | 6.26 |
Both models were timed in the same run on the same GPU. Compared with ESM2-8M in bf16, the deployable setting for both, flashEE is 2.2× faster and uses 2.7× less memory per sequence.
Most common quantisation formats require a dequantisation step to make the weights available for calculations.
For native calculations, FP32 is the most common high-precision, high-resolution format. However, these large numbers increase the attention matrix size. Alternative formats include BF16 (considered the gold standard) and FP8 and FP4 (MXFP4) on NVIDIA Blackwell GPUs, which use specially optimised cores, although I did not have access to these during this project.
| Format (B=32) | Weights MiB | ms/seq | vs fp32 | Peak MiB |
|---|---|---|---|---|
| fp32 | 18.89 | 1.126±0.009 | 1.00× | 179.0 |
| bf16 | 9.46 | 0.325±0.002 | 3.47× | 114.4 |
| int8, packed | 4.85 | 1.136±0.014 | 0.99× | 185.7 |
| int4, packed | 2.51 | 1.150±0.014 | 0.98× | 183.4 |
| int2, packed | 1.33 | 1.149±0.014 | 0.98× | 182.2 |
While quantised weights greatly reduce the memory footprint, they do not speed up inference or reduce peak memory usage at this scale. Packed 4-bit weights are 2% slower than fp32 and use more peak memory than bf16 (183.4 vs 114.4 MiB) because the dequantised copy that the arithmetic actually uses is still in fp32, so the model holds both at once, whereas bf16 weights do not require that process.
Model quality varies depending on the quantisation protocol used. I evaluated int quantisations and my own implementations of the formats llama.cpp uses. Additionally, since my model is heavily normalised, I examined how nf4, a format designed to preserve normalised values, split into Gaussian quantiles and not uniformly, compares with other formats within the same bit class.
| Format | bits/weight | ppl | top-1 | ΔNLL retained |
|---|---|---|---|---|
| fp32 | 32.00 | 10.372±0.012 | 29.79 | 100% |
| q6_k | 6.56 | 10.379±0.013 | 29.78 | 99.8% |
| q5_k_s | 5.50 | 10.403±0.012 | 29.71 | 99.2% |
| q4_1 | 5.00 | 10.507±0.008 | 29.42 | 96.8% |
| nf4:g32 | 4.50 | 10.535±0.008 | 29.34 | 97.2% |
| int4:g32 / q4_0 | 4.50 | 10.652±0.008 | 29.03 | 94.7% |
| nf4 | 4.04 | 10.672±0.022 | 28.94 | 93.5% |
| e2m1 (MXFP4) | 4.04 | 10.727±0.029 | 28.81 | 92.3% |
| int4 (per-channel) | 4.04 | 11.030±0.039 | 28.06 | 85.2% |
| q3_k_s | 3.44 | 10.818±0.031 | 28.56 | 87.8% |
| Int3 (per-channel) | 3.04 | 15.643±1.486 | 19.68 | 38.2% |
ESM2-8M quantised using the same protocol shows a much larger performance drop, where at int4, it retains only 68.7% of its full-precision ProteinGym score, while flashEE preserves 91.1% ± 4.0. As a result, ESM2-8M, which leads by 18.6% at full precision, falls behind flashEE by 33% ± 2.5 on eukaryotic assays upon 4-bit quantisation.
| 4 bit quantisation | flashEE | ESM2-8M |
|---|---|---|
| int4 | 91.1% ± 4.0 | 68.7% |
| e2m1 | 91.7% ± 4.1 | 79.7% |
| nf4 | 94.3% ± 2.5 | 82.9% |
The quantisers in this project were externally validated; nf4 matches bitsandbytes to 7.6e-08, and asymmetric grouped int4 matches weights recovered through PyTorch's tinygemm kernel to 4.7e-03.
Activations proved far more fragile than weights: 8-bit costs nothing measurable, but 4-bit collapses to 37.5% of the baseline
Quantisation-aware training (QAT), run for a further 20k steps with the quantiser in the loop, closes the four-bit gap almost entirely. QAT-nf4 and QAT-e2m1 reach FP32 parity, whereas QAT-int4 improves but does not quite close it.
| Format | top-1 | vs fp32 | ppl | BLOSUM62 |
|---|---|---|---|---|
| fp32 | 29.790 | NA | 10.372 | 0.4968 |
| PTQ, int4 | 28.060 | −5.81% | 11.030 | 0.4720 |
| PTQ, e2m1 | 28.810 | −3.29% | 10.727 | 0.4780 |
| PTQ, nf4 | 28.940 | −2.85% | 10.672 | 0.4889 |
| QAT, int4 | 29.583±0.041 | −0.69% | 10.442 | 0.4908 |
| QAT, nf4 | 29.851±0.044 | +0.20% | 10.340 | 0.5042 |
| QAT, e2m1 | 29.860±0.068 | +0.23% | 10.346 | 0.5038 |
Sparse autoencoder (SAE) analysis helped to localise the damage. At 4 bits, roughly half of flashEE's features remain functional, compared with only one-eighth under random perturbation of the same magnitude. In contrast, isotropic noise spreads into low-variance directions and overwhelms them. To determine whether this extends beyond my own model, I repeated the analysis using InterPLM's published dictionaries for ESM2-8M and ESM2-650M, each with 10,240 features, trained and preannotated, across 300 reviewed eukaryotic Swiss-Prot proteins (121,738 residues). A matched-noise control perturbed weights to the same norm, with a 2,000-resample bootstrap over proteins. Like prior quantisation runs, int8 was lossless, while ESM2-8M uniform int4 retained 83% of full-precision F1, compared with 94% for e2m1 and 96% for nf4, matching the near-zero level-density pattern observed in flashEE. The size effect held too, with ESM2-8M's degradation resembling flashEE's; its signal-peptide detector retained 69% against flashEE's 86%, while ESM2-650M was unaffected by every 4-bit format tested, retaining full-precision accuracy, supporting the idea that smaller models are more prone to degraded representations than large ones.
CPU inference is another topic I decided to investigate; it could let more people run it, and some packages are better optimised for specialised formats like int8-dynamic using fbgemm on x86 hardware. Although most people wouldn't want to do this with PLMs, int8-dynamic performs multiplication using quantised weights; still, those results sit inside the noise: the run-to-run spread at each batch size is 3-12%, larger than every ratio in the table.
| B= | fp32 ms/seq | int8-dynamic | ratio |
|---|---|---|---|
| 8 | 73.13 | 70.37 | 1.04× |
| 32 | 77.67 | 74.51 | 1.04× |
| 64 | 87.47 | 88.17 | 0.99× |
| 128 | 96.07 | 93.47 | 1.03× |
This project successfully created a specialist model that runs faster than the closest general-purpose model of comparable size. This showed that even the smallest protein language models can be improved by borrowing methods from NLP without degrading representation. On top of that, the quantisation arc showed that these small models can be quantised to take up less space while still encoding the "truth" just as well, even though the savings are in storage rather than speed or peak memory.
This model was trained only on eukaryotic sequences and uses a different architecture than ESM2-8M. Whether their architecture was better for this dataset, or mine for theirs, was not tested.
This project could take many different paths; I have only scratched the surface. Firstly, this is a tiny model. Time and compute constraints limited me, and each iteration took ~8 days to run. A larger, specialist model (~50M parameters) would be much better and could scale much more widely, and much more likely deployed by someone as part of their pipeline. Another constraint is data: while extensive, the UniProt library could be supplemented with other databases, potentially by filtering databases used by other organisations and ensuring we preserve only the high-quality ones.
One direction I was looking into was distillation (specifically from the ESMC-300M model), but it would require a custom adapter (since my model encodes the amino acids slightly differently in space) and would take much more time and compute than I had at hand, most likely months.
Speaking of architecture tweaks, I was also considering implementing a mixture-of-experts (MoE) architecture. It would allow only parts of the model to activate in response to a specific task, effectively reducing the compute required to complete it. While it works well in practice, I found it unnecessary for my model, since the weights and their uses overlapped too much at this scale.
Another thing I wanted to borrow from NLP systems was gated attention. Systems like that are built on SDPA attention and check which attention is relevant and which isn't; it's most famously used in the Qwen models, not really used in sequence-only PLMs. Even so, at this scale it would cost more than it would save.
Multiple sequence alignment (MSA), while not necessary, has also been shown to improve zero-shot prediction frequencies. Despite that, I opted not to include it since MSAs can be quite tricky to get right, and misaligned MSAs can lead to consistently degraded representations. That said, I think the approaches taken by the people behind the MSA Pairformer (https://www.biorxiv.org/content/10.1101/2025.08.02.668173v1), which I came across toward the end of my project, have the best potential to be fast, efficient PLMs that go against the grain on model scaling. If I had more time/compute, that would probably be my top priority to implement next.
main.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)
sae.py sparse autoencoder + feature survival under quantisation
validate_external.py our quantisers vs bitsandbytes and ggml
report.py leaderboard + champion selection
plots/ every figure in make_figures.sh
flashee/ pip-installable inference wrapper (torch only, no training stack)
Maciej Szczesny, as part of Stracquadanio Lab, 2026. This project was generously supported by CCPBioSim and EPSRC.