🧬 biology

Task- and dataset-specific information in protein language models

This study reveals that the most informative embeddings for protein language model downstream tasks are often found in intermediate layers rather than the final layer, with the optimal layer depending on whether the task involves individual residues, whole proteins, or artificial sequences.

Original authors: Roman Joeres, Ilya Senatorov, Olga V. Kalinina

Published 2026-08-13

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Roman Joeres, Ilya Senatorov, Olga V. Kalinina

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ⚕️ This is an AI-generated explanation of a preprint that has not been peer-reviewed. It is not medical advice. Do not make health decisions based on this content. Read full disclaimer

Imagine you have a massive, magical library where every book is a recipe for building a living machine called a protein. For years, scientists have been trying to teach computers to read these recipes and understand what the machines do. To do this, they built "Protein Language Models" (PLMs). Think of these models as super-smart students who have read billions of these protein recipes. They don't just memorize the words; they learn the hidden grammar and style of the language. Once they've studied, these students can turn a protein's recipe into a secret code (an "embedding") that other computers can use to predict things like: Will this protein glow? Will it stick to a virus? Will it dissolve in water?

For a long time, everyone assumed that the "smartest" part of these student models was at the very end of their brain—the final layer. It was like assuming that after reading a whole book, the last sentence you read holds the most important meaning. But what if the most useful clues were actually hidden in the middle of the book, or even in the first few pages? This paper asks a simple but tricky question: Where in these giant AI brains does the real magic happen? Do we really need to look at the very last layer to get the best answers, or are we missing something important by ignoring the layers in between?

The researchers behind this study decided to play detective. They took 13 different protein-learning models and tested them on 15 different tasks, ranging from predicting how stable a protein is to figuring out where it lives inside a cell. They didn't just look at the final answer; they peeked inside the models at every single step of their thinking process, from the first layer to the last.

Here is what they discovered, and it turns out the old rule of "the last layer is the best" is mostly wrong.

The Middle is Often the Sweet Spot
When the scientists tested these models on tasks involving whole proteins (like predicting if a protein is soluble or where it lives in a cell), they found that the very last layer was rarely the winner. Instead, the "best" layer was usually somewhere in the middle—often between the 10th and 90th percentiles of the model's depth. It's as if the student reads the book, understands the main plot in the middle chapters, and then starts to get confused or over-analyze things by the time they reach the final page. In fact, for many tasks, the performance actually dropped in the deepest layers.

It Depends on What You're Asking
The paper found that the "best" layer isn't the same for every job. It depends heavily on the type of data you are using:

The "Mutant" Datasets: When the data came from "Deep Mutational Scanning" (where scientists take one specific protein and make thousands of tiny, slightly different versions of it), the models worked best using the shallow, early layers. It seems the early layers are great at spotting small, local details, like how changing one letter in a word changes its meaning.
The "Diverse" Datasets: When the data came from a huge mix of many different, unrelated proteins (like a library of thousands of different books), the models worked best using the deeper layers. Here, the model needs to understand big, general rules that apply to many different proteins, which takes more "thinking" and deeper layers to figure out.

The "Artificial" Problem
One of the most surprising findings was about "artificial" proteins. The researchers tested the models on proteins that were designed by computers rather than found in nature. The models struggled significantly with these. Even if the artificial proteins were functional, the models couldn't predict their properties well. This suggests that these AI models have learned to understand the "language" of nature's evolution, but they haven't quite figured out the "language" of computer-designed proteins. They are fluent in the old dialect but stumble over the new slang.

Why the Last Layer Isn't Always King
The paper explains that this happens because of how these models are trained. They are first trained to guess missing words in a sentence (a task focused on small, local details). When you ask them to do a big-picture task (like predicting the stability of a whole protein), the very last layer is still trying to do that small-word guessing job, which gets in the way. However, if you "fine-tune" the model specifically for a big-picture task, then the layers start to work better in order, with the last layer becoming the most useful again. But in the standard, pre-trained state, the middle layers often hold the gold.

The Takeaway
This study suggests that scientists shouldn't just blindly grab the final layer of a protein AI model for their work. Instead, they should look deeper into the model's "brain" to find the specific layer that matches their task. If they are looking at tiny mutations, they should look at the early layers. If they are looking at a diverse mix of proteins, they should look deeper. And if they are working with computer-designed proteins, they should be extra careful, because the models might not understand them as well as they understand nature's own creations.

In short, the paper reveals that these AI models are complex, layered thinkers, and the "smartest" part of their brain changes depending on what question you ask them. The last page of the book isn't always the most important one; sometimes, the best answer is right in the middle.

Technical Summary: Task- and Dataset-Specific Information in Protein Language Models

Problem Statement
Protein Language Models (PLMs) have become central to computational biology, translating amino acid sequences into latent-space embeddings for diverse downstream tasks (DTs). A prevailing consensus in the field assumes that the embeddings from the final layer of a PLM are the most informative, a heuristic borrowed from deep learning models trained end-to-end on specific tasks (e.g., ResNet-50 on ImageNet). However, PLMs are typically pre-trained via self-supervised objectives (e.g., Masked Language Modeling or Next Token Prediction) on vast corpora of protein sequences and then used as feature extractors for downstream tasks with different objectives. It remains unclear whether the "last layer is best" assumption holds for PLMs, how information is distributed across intermediate layers, and how factors such as dataset composition (natural vs. artificial proteins) and task granularity (residue-level vs. protein-level) influence the optimal embedding layer.

Methodology
The authors conducted a systematic analysis across 13 PLMs from five model families (ESM-2, ESMC, ProtT5, ProstT5, ProGen2, ProtGPT2) spanning four different architectures (BERT, Transformer, T5, GPT). These models were evaluated on 15 downstream tasks derived from 11 distinct datasets, covering both residue-level and protein-level predictions.

The methodology employed two primary approaches:

Linear Probing: Probe models (linear classifiers and $k$ -nearest neighbors) were trained on embeddings from every layer of each PLM to evaluate performance on specific DTs.
Latent Space Analysis: The authors computed three metrics to characterize the geometry and complexity of the latent spaces at each layer:
- Intrinsic Dimension (ID): Estimated using the TwoNN approach to determine the number of independent variables required to describe the space.
- Neighborhood Overlap: Measured the similarity of neighbor sets between consecutive layers.
- Variance@10: The fraction of variance explained by the first 10 principal components.

Additionally, the study included ablation studies involving dataset subsampling to determine data efficiency in identifying optimal layers and experiments comparing natural proteins against artificially generated sequences (via Rosetta and ProGen).

Key Results

The Last Layer is Rarely Optimal: Contrary to common practice, the embeddings from the final layer of a PLM rarely yielded the best performance. In only 17.92% of cases did the last layer outperform all others. For most protein-level tasks, performance peaked in intermediate layers (between the 10th and 90th percentiles of model depth) and declined in the deepest layers.
Task and Dataset Dependence: The location of the most informative layer is heavily dependent on the dataset structure:
- Deep Mutational Scanning (DMS) Datasets: For datasets containing variants of single proteins (e.g., Fluorescence/GFP, GB1), embeddings from shallow layers performed best.
- Diverse Natural Protein Datasets: For datasets comprising many distinct natural proteins (e.g., DeepLoc2.0, DeepSol, SCOPe40), performance generally increased with layer depth, peaking in deeper layers before a slight decline at the very end.
Residue vs. Protein Level: Residue-level tasks (e.g., secondary structure prediction, binding site identification) consistently showed a steady increase in performance with layer depth, aligning with the residue-level nature of pre-training objectives. Protein-level tasks exhibited non-monotonic performance curves.
Pre-training vs. Downstream Alignment: The performance trends correlate with the alignment between the pre-training objective and the downstream task. When the objectives align (e.g., residue-level tasks), performance improves monotonically. When they diverge (e.g., applying a residue-trained model to whole-protein tasks), the "drift" in deeper layers causes performance to drop. Fine-tuning a PLM on a specific task restores monotonic improvement across layers for that task.
Artificial Proteins: PLMs demonstrated significantly reduced performance when predicting properties of artificially generated proteins (e.g., Rosetta-designed miniproteins or ProGen-generated lysozymes) compared to natural proteins or their mutants. This suggests PLMs encode representations of naturally evolved sequence spaces rather than abstract, universal protein functions.
Data Efficiency: The study found that 15–20% of the training data is sufficient to identify a layer that achieves close-to-optimal performance (≥95% of the best), indicating robustness in layer selection.

Significance and Claims
The authors claim this work provides the first comprehensive study of PLM layer behavior across architectures and tasks, challenging the standard practice of using only the final layer embeddings. The paper posits that:

Information is Layer-Specific: PLMs do not simply accumulate information linearly; early layers capture fine-grained, local details (relevant for DMS/variants), while deeper layers capture general, global rules (relevant for diverse natural proteins).
Dataset Structure Dictates Utility: The utility of a specific layer is defined more by the dataset's diversity (single-protein variants vs. diverse proteomes) than by the task type alone.
Limitations on Artificial Sequences: PLMs struggle to generalize to artificial protein designs, indicating their learned representations are specific to the statistical properties of natural evolution rather than a generalizable understanding of protein physics or function.

The paper concludes that selecting the optimal embedding layer requires careful consideration of the dataset and task, and that the "last layer" heuristic is often suboptimal for protein language modeling.

Technical Summary: Task- and Dataset-Specific Information in Protein Language Models

More like this