PLMSommelier

The last layer of a protein language model is almost never the best one. This tool finds the layer that is, and hands you back a truncated model.

Implements the tool described in Task- and dataset-specific information in protein language models, which probed 13 PLMs across 15 downstream tasks and found the deepest layer won in only 17.9% of cases.

Install

uv tool install plmsommelier   # isolated, no venv to manage yourself
# or: pipx install plmsommelier
# or: pip install plmsommelier   # into a virtualenv

# CPU-only (much smaller download; also what you want on a machine with no GPU)
pip install torch --index-url https://download.pytorch.org/whl/cpu
pip install plmsommelier

# GPU: pick the wheel matching your driver, then install as usual
pip install torch --index-url https://download.pytorch.org/whl/cu126   # CUDA 12
pip install torch --index-url https://download.pytorch.org/whl/cu130   # CUDA 13
pip install plmsommelier

Quickstart

examples/fluorescence_sample.csv ships in the repo (a 500-row subsample of the TAPE fluorescence benchmark -- see examples/README.md), so this runs in about a minute on CPU with no data of your own needed:

plmsommelier suggest examples/fluorescence_sample.csv facebook/esm2_t6_8M_UR50D \
    --task regression --out ./my-esm-truncated

data and model are positional (DATA MODEL); --data/--model work identically if you prefer named flags.

model    facebook/esm2_t6_8M_UR50D
dataset  fluorescence_sample  (regression, pearson, knn probe)
data     400 train / 100 val

best layer      0 of 6 (0% depth)   pearson = 0.4238
last layer      6               pearson = 0.2369
gain over last  +78.9%
seed agreement  100% (layers chosen across seeds: [0])
confidence      moderate   (plateau agreement 100%, separation 1.51x, seed spread 0% of depth)

layer performance (pearson, +/- 1 sd across 5 seeds):
    0 +0.4238 ########################################-----   +/-0.0537  <- best
    1 +0.3764 ####################################---------   +/-0.1012
    2 +0.3770 ####################################-----   +/-0.0634
    3 +0.1572 ###############------   +/-0.0736
    4 +0.0432 ####----   +/-0.0519
    5 +0.2014 ###################-------   +/-0.0824
    6 +0.2369 #######################------   +/-0.0710  (last)

(Layer 0 is the token embedding output before any transformer block -- a real, if unglamorous, answer on this tiny 6-layer model and 500-row sample; larger models and datasets typically pick somewhere mid-stack. See the paper for what "best layer" looks like at scale.)

./my-esm-truncated is a normal HuggingFace model directory holding only the chosen number of blocks. It loads anywhere the original did, runs faster, and scores at least as well on your task:

from transformers import AutoModel, AutoTokenizer

model = AutoModel.from_pretrained("./my-esm-truncated")
tok = AutoTokenizer.from_pretrained("./my-esm-truncated")

T5-encoder checkpoints (Ankh, ProtT5, ProstT5, ...) load via AutoModelForTextEncoding instead of AutoModel. You don't need to remember which is which: the exact loading snippet for whatever model you just truncated is always written into the generated {out}/README.md.

Any HuggingFace protein language model works — just pass its id to model. Checkpoints that ship custom modeling code need --trust-remote-code. See Extending to a custom PLM below for models that need more than that.

Python API

The CLI is a thin wrapper around five importable functions (plmsommelier/__init__.py), wired together in exactly this order by cli.py::suggest_layer if you'd rather read one function than five:

from plmsommelier import embed_layers, load_dataset, load_model, save_truncated, select_layer

plm = load_model("facebook/esm2_t6_8M_UR50D")
ds = load_dataset("examples/fluorescence_sample.csv", task="regression")

train_embeddings = embed_layers(plm, ds.train_seqs)
val_embeddings = embed_layers(plm, ds.val_seqs)

result = select_layer(ds, train_embeddings, val_embeddings, model_name=plm.model_id)
print(result.summary())  # the block printed by the CLI
print(result.curve_plot())  # the ASCII per-layer curve
result.to_dict()  # what --json writes out

save_truncated(plm, result, "./my-esm-truncated")

Result (plmsommelier.select.Result) carries the full per-layer curve (result.curve), the per-seed scores behind the stability check (result.seed_curves, result.layer_sigma), and the confidence verdict -- see Confidence and what to do about it.

Input format

A CSV with sequence and a label column:

--task (regression, classification or multi-label) is inferred from the label column when that's unambiguous: multiple label columns -> multi-label, non-numeric or exactly two distinct values -> classification, a float dtype -> regression. A whole-number column with more than two values could be class codes or an integer-valued regression target, and inference refuses to guess there -- pass --task explicitly.

--max-seq-len (default 2000 residues) drops longer sequences before splitting or subsampling, so a single outlier protein can't end up alone in an unbounded batch -- attention memory is quadratic in length. Pass 0 to disable it.

Device is auto-detected -- CUDA, else Apple Silicon's mps, else CPU -- and can be forced with --device cpu / --device cuda / --device mps. A visible CUDA device only counts if the installed torch build actually ships kernels for it: an older card (e.g. anything pre-Turing) against a build that dropped support for it would otherwise crash on the first forward pass with no kernel image is available for execution on the device instead of just running on the CPU. Auto-detect falls through with a warning in that case; pass --device cuda explicitly to force it anyway.

CLI reference

plmsommelier suggest --help is always authoritative -- every flag is generated directly from suggest_layer's Python signature and docstring, so it can't drift out of sync with this table.

Flags shared by every command, handled by the top-level launcher:

Extending to a custom PLM

There is no model registry -- model takes any HuggingFace id or local path, and works with zero configuration for most architectures. Reach for the options below only as far as your checkpoint actually needs, in order of how rare each one is:

1. Nothing. load_model (plmsommelier/model.py) is architecture-agnostic: it finds the transformer blocks and final norm by attribute name, detects the sequence-length ceiling from the tokenizer/config/causal-mask buffer, and picks the right Auto* loading class. Most encoder PLMs on the Hub just work.

2. --trust-remote-code for checkpoints that ship custom modeling code in their own repo (an auto_map in config.json). Same flag transformers itself uses.

3. A pretraining-convention entry in _QUIRKS (plmsommelier/model.py), keyed on config.model_type, for models whose config doesn't record how they were actually pretrained -- HuggingFace has no field for "residues are space-separated tokens", for instance. Supported keys:

   key	effect
   space_join	join residues with spaces before tokenizing (ProtTrans-style)
   residue_map	a str.translate table applied to the sequence first (e.g. rare residues -> X)
   prefix_text	text prepended before tokenizing (e.g. a mode/direction token)
   dtype	force this torch dtype instead of the auto-detected default
   encoder_only	force AutoModelForTextEncoding (True) vs. AutoModel/AutoModelForCausalLM (False), overriding the is_encoder_decoder config check

   T5/BERT/ALBERT (the ProtTrans family) are the worked example already in the table:

   "t5": dict(
       space_join=True,
       residue_map=str.maketrans("UZOB", "XXXX"),
       dtype=torch.bfloat16,   # T5 activations overflow in fp16
       encoder_only=True,      # loads T5EncoderModel; the decoder is never built
   ),
   

   Add a family the same way: an entry keyed on its model_type (AutoConfig.from_pretrained(your_model).model_type tells you the key), with only the keys it actually needs.

4. A _REGISTRY_PACKAGES entry for architectures that live in a third-party package rather than in transformers itself -- these checkpoints carry no auto_map, so --trust-remote-code can't reach them; the package must be imported first so it registers itself with the Auto* classes. multimolecule/proteinbert is the current example: install plmsommelier[multimolecule], and model.py imports multimolecule before loading whenever config.model_type == "proteinbert". Add a new package the same way: "model_type": ("import_name", "pip_extra_name"), plus an extra in pyproject.toml if it isn't already installed alongside plmsommelier.

5. No fork needed for a one-off. _QUIRKS and _REGISTRY_PACKAGES are plain module-level dicts -- from your own script, import plmsommelier.model as m; m._QUIRKS["your_model_type"] = {...} before calling load_model works without touching this repo. Treat that as a private escape hatch, though (the leading underscore is deliberate) -- if it's a real, reusable family, a PR adding it to the table is the better home for it.

Whatever route you take, verify a new family the way tests/test_model.py verifies the built-in ones: the truncation round-trip (save_truncated then reload reproduces the same layer's output), sequence-length detection not silently truncating valid input, and the final-norm invariant -- see CONTRIBUTING.md's invariants list for what each of those actually guards against.

Tested models

Checkpoints below were loaded and probed end-to-end (load_model + embed_layers) against the real weights, not just read from the code — either by the repo's own real-weight test fixtures (tests/test_model.py) or by a manual smoke test. Anything not listed still has a good chance of working — load_model is architecture-agnostic — it just hasn't been verified here yet.

ProtBert, ProtAlbert, ProtT5, Ankh, ProstT5, and ProtGPT2 are handled by name in _QUIRKS but have no real-weight test coverage yet — treat them as likely-to-work, not confirmed.

Confidence and what to do about it

Every run redraws the training data (and, where there's enough of it, the validation data) --n-seeds times and re-scores every layer, to check whether the chosen layer survives resampling. That feeds a confidence verdict -- high, moderate, or low -- printed as part of the summary and recorded in --json output. unmeasured means the check wasn't run at all (--n-seeds 0), which is different from low: low means "we checked, and the pick isn't stable"; unmeasured means "we didn't check".

The verdict is the worst of three independent signals, not an average, so one strong number can't paper over another weak one:

• seed agreement -- how often resampling lands back inside the same plateau (the tolerance-band region of tied-best layers), not just on the exact same layer. Two seeds picking neighbouring, statistically tied layers count as agreement, not disagreement.
• peak margin -- how many (paired) standard deviations separate the plateau from the best layer outside it, across resamples. This is a resampling-stability statistic, not a hypothesis test: seeds are overlapping subsamples of one dataset, not independent draws, so no p-value is implied.
• seed spread -- how far, on average, resampled picks land from the chosen layer, as a fraction of the network's depth.

A low verdict prints a warning block to stderr (not stdout, so --json and piped output stay clean) with concrete, run-specific hints. The most important thing to know before reading them: raising --n-seeds measures confidence more precisely -- it does not raise it. Only one of the following actually reduces variance:

1. Raise --sample, if the hint says there's headroom (the tool tracks how many rows were available versus used). Score variance falls roughly as 1/N, so ~3x the data tightens the estimate by roughly 1.7x. If you're already using every row in the file, this can't help -- the hint says so explicitly rather than telling you to raise a flag that won't do anything.
2. Try --probe lr (or raise --k for the kNN probe). A linear probe has much lower variance than kNN at typical sample sizes; two probes agreeing on a layer is stronger evidence than either one's confidence number alone.
3. Check the validation set. Below roughly 50 val rows, none of these statistics are trustworthy -- --sample and the train/val split both drive this.
4. Check the data for heavy class imbalance, a near-constant label column, or a split column that leaves a tiny validation side.
5. Accept the plateau. If the curve is genuinely flat, low confidence is the correct answer, not a problem to fix -- the tool already returns the shallowest layer in the tied region, which is the smallest and fastest model at equal performance.

Development

See CONTRIBUTING.md for the dev install and workflow, and RELEASING.md for how a release is cut and published.

pip install -e ".[dev]"
pytest tests -q                      # full suite
pytest tests -q -m "not weights"     # no model downloads needed

Citation

@article{joeres2026taskdataset,
  title  = {Task- and dataset-specific information in protein language models},
  author = {Joeres, Roman and Senatorov, Ilya and Kolchina, Anastasia and
            Klakow, Dietrich and Kalinina, Olga V.},
  journal = {arXiv preprint arXiv:2608.12090},
  year   = {2026}
}

See CITATION.cff for the software-citation form.

License

MIT — see LICENSE.