nanoprot-mamba-XS
A protein language model (mamba architecture) — 9.2M (9,186,688) parameters, trained on 110.2M (110,240,256) UniRef50 residues.
nanoprot-mamba-XS is part of the nanoprot suite: a Pythia-style matrix of protein
language models spanning three architectures (gpt2, esm2, mamba) and four
scales (XS/S/M/L), each trained from scratch on UniRef50 under a matched,
Chinchilla-style data budget. The suite is built for controlled comparison —
same data, same tokenizer, one variable at a time.
Headline result
Validation bits-per-residue (lower is better): 4.0191 ± 0.0007 (n=3 seeds)
Directly comparable to other autoregressive nanoprot models (
gpt2,mamba) — same 33-token vocabulary, same AR objective, same data budget. Not comparable to theesm2(masked-LM) models, whose metric is a different quantity.
Model details
| Architecture | mamba |
| Objective | AR (causal-language-modeling) |
| Scale rung | XS |
| Parameters | 9.2M (9,186,688) |
| Layers (depth) | 7 |
| Hidden size (d_model) | 448 |
| Attention heads | 1 |
| Max sequence length | 512 |
| Vocabulary | 33-token residue alphabet (ESM-2) |
| d_state | 16 |
| d_conv | 4 |
| expand | 2 |
| dt_rank | 28 |
| Precision | bf16 |
Training
| Data | UniRef50 release 2026_01 (28-Jan-2026), 60,251,814 sequences; held-out final shard for validation |
| Tokenizer | esm2 — 33-token residue alphabet (shared across the whole suite) |
| Optimizer | Muon (matrices) + AdamW (embeddings/scalars), weight_decay=0.1 |
| Batch size | 524,288 residues/step |
| Optimizer steps | 210 |
| Residues seen | 110.2M (110,240,256) |
| Param/data ratio | 12.0 (Chinchilla-style) |
| Total FLOPs | 3.997e+16 |
| Wall-clock | 0.03 h (2 min) on 4 GPU(s) |
| Seed | 0 (siblings: see below) |
| nanoprot version | 0.5.0 |
Evaluation
Evaluated on a held-out UniRef50 shard. Validation bits-per-residue (lower is better): 4.0191 ± 0.0007 (n=3 seeds).
Directly comparable to other autoregressive nanoprot models (gpt2, mamba) — same 33-token vocabulary, same AR objective, same data budget. Not comparable to the esm2 (masked-LM) models, whose metric is a different quantity.
Intended use & limitations
Research use: learning protein representations, extracting residual-stream features, mechanistic-interpretability probing, and architecture comparison. Trained only on UniRef50 sequences — not for clinical or diagnostic use, and not aligned to any downstream task out of the box.
How to load
Install nanoprot (pip install nanoprot), download this repo, and point the
arch-aware loader at the folder — it works for any nanoprot architecture
(gpt2 / esm2 / mamba), reading the embedded config and selecting the right
tokenizer automatically.
from nanoprot.training.checkpoint import load_pretrained
model, cfg, meta, tokenizer = load_pretrained(
"path/to/this/repo", device="cpu", return_tokenizer=True,
)
model.eval()
# meta carries the trained-artifact facts (params, FLOPs, val metric, ...)
The nanoprot suite
Hub: yagizdevre/nanoprot-mamba-XS (seed 0 is the default; siblings on branches seed1,
seed2). This model's sibling seeds:
nanoprot-mamba-XS-s0— final bits-per-residue 4.0196nanoprot-mamba-XS-s1— final bits-per-residue 4.0184nanoprot-mamba-XS-s2— final bits-per-residue 4.0193
The full suite spans {gpt2, esm2, mamba} x {XS, S, M, L} x {seed 0,1,2}.
See the nanoprot repository for the
complete grid and the scaling-curve comparisons.
Citation
@software{nanoprot,
author = {Devre, H. Yagiz},
title = {nanoprot: a minimal training framework for protein language models},
year = {2026},
url = {https://github.com/ygzdvr/nanoprot}
}
Reproducibility
Trained with nanoprot v0.5.0 (corpus prepared 2026-06-01T12:51:52+00:00). The exact, complete training
config is in config.yaml (also embedded in meta_000210.json).
Re-train with:
python -m scripts.train --config config.yaml