Task- and dataset-specific information in protein language models
Roman Joeres, Ilya Senatorov, Anastasia Kolchina, Dietrich Klakow, Olga V. Kalinina
cs.LG, q-bio.BM
2026-08-12
Probing every layer of 13 protein language models on 15 tasks: last-layer embeddings win only 17.9% of the time; dataset makeup, not the task, decides the best layer.
The default recipe for using a protein language model (PLM) is to take embeddings from the last layer and train a prediction head on top. The habit comes from computer vision, where a network trained end to end really does improve with depth: every layer of an ImageNet classifier serves the final linear head. PLMs are pretrained self-supervised on a different objective, usually guessing masked amino acids, so nothing guarantees the last layer is the most useful one for predicting solubility or subcellular location. Earlier work had poked at single models; this paper runs the first systematic sweep, 13 PLMs (five ESM-2 sizes, two ESMC, ProtT5, ProstT5, three ProGen2 sizes, ProtGPT2, up to 3B parameters) across 15 downstream tasks on 11 datasets, probing every layer.
The core instrument is the linear probe: freeze the PLM, train a small linear head on each layer's output, and read probe accuracy as a measure of how much task-relevant information that layer holds. Protein-level tasks average residue vectors with mean pooling; residue-level tasks (secondary structure, ligand binding) use per-position embeddings. Three geometry metrics computed per layer, intrinsic dimension via TwoNN, neighborhood overlap between adjacent layers, and variance explained by the first 10 principal components, track how the latent space evolves with depth.
Two follow-up experiments pin down the cause. First, the authors evaluate each layer on the pretraining objective itself, masked-language-model loss for ESM-2 and next-token prediction for ProGen2. Second, they fine-tune ESM-2 150M on six downstream tasks and re-probe every layer, checking whether the non-monotonic curves straighten out and what the fine-tuning costs on the original objective.
The last layer is the best choice in only 17.92% of model-task combinations. Protein-level tasks climb over the first layers, peak somewhere between the 10th and 90th percentile of depth, and drop again near the end. Residue-level tasks rise monotonically.
On the pretraining objective itself, both families improve almost monotonically with depth. After fine-tuning, the downstream curves become monotonic too. The non-monotonic behavior is an artifact of two-stage training with mismatched objectives, not a defect of the models.
Which layer wins is set by the dataset, not the task. Two tasks on the same dataset produce nearly identical layer curves: correlation 0.978 between the two DeepLoc2.0 tasks, 0.996 between 3-class and 8-class secondary structure on SCOPe40, and 0.681 on the fluorescence pair.
| Dataset type | Examples | Mean sequence identity | Where the best layer sits |
| DMS variant libraries | Fluorescence (54,024 GFP variants), GB1 (149,361 variants) | 88% to 97% | Early layers, deep layers fall off |
| Diverse natural proteins | DeepLoc2.0, DeepSol, SCOPe40 | 13% to 17% | Rises with depth, sags only in the deepest 5% |
Early layers capture local detail and mutational context; later layers capture general rules that separate proteins, mirroring the syntax-early, semantics-late pattern known from text LLMs. On individual tasks the ESMC models' best layer beats their last layer by 165% to 170%. On fluorescence regression, best-layer probe scores range from 0.58 (ESM-2 8M) to 0.74 (ProGen2-large).
On artificial proteins the models largely fail. Rosetta-designed proteins in the Rocklin stability dataset score far worse than the DMS and natural portions. On 90 ProGen-generated lysozymes, activity prediction reaches a Spearman correlation of about 0.2, against about 0.6 for natural sequences, even though the artificial activities span a wider range.
Practically, 15% to 20% of the training data is enough to find a layer reaching at least 95% of the best achievable performance, and larger models tolerate sparser data better.
If you embed proteins with ESM today, stop defaulting to the last layer. A cheap layer sweep on a fraction of your labels buys real performance. The deeper lesson matters for protein design: PLMs encode the structure of naturally evolved sequence space, not protein function in general. Ranking the stability or activity of freshly designed proteins with a PLM is unreliable, and this holds for physics-based Rosetta designs as much as for PLM-generated ones. Any workflow that designs proteins and scores them with a PLM should carry this caveat.
The authors cap the study at 3B parameters, citing storage (about 10 TB of protein embeddings alone) and hardware; residue-level experiments skip ESM-2 3B, ProGen2-large and ProtGPT2. ProtGPT2 uses byte-pair tokenization, so one token covers several amino acids and residue-level tasks are impossible for it. The artificial-protein conclusion rests on two cases, stability and lysozyme activity, and the authors extend it only cautiously. Probing measures linear separability, a lower bound on information, so deeper layers may hold more than the probes show. Finding the best layer still requires labeled data; no unsupervised recipe is offered.