Task- and dataset-specific information in protein language models

This study investigates the internal representation of protein language models (PLMs) by analyzing 13 different models across 15 downstream tasks (DTs). While it is common practice to use the final layer's embeddings for downstream applications, the authors hypothesized that this assumption—borrowed from end-to-end trained computer vision models—may not hold for PLMs trained via self-supervised objectives. To test this, the researchers used linear probes to evaluate the informativeness of every layer in each model, complemented by latent-space analysis metrics such as intrinsic dimension and neighborhood overlap.

The researchers found that the final layer is rarely the most effective for downstream tasks. For protein-level tasks, performance typically peaks in intermediate layers (between the 10th and 90th percentiles of model depth) and often declines toward the final layer. This pattern is linked to the pre-training objective: models are optimized to predict masked tokens, causing the final layer to revert to representations similar to the input, which may be less useful for global, sequence-level protein properties.

Furthermore, the study identifies a clear distinction between task types. For residue-level tasks, performance tends to improve monotonically with depth. For whole-protein tasks, the optimal layer is heavily influenced by the dataset: deep mutational scanning (DMS) datasets favor shallow layers, whereas diverse natural protein datasets favor deeper layers. The authors also observed a significant performance drop when applying these models to artificially generated proteins, suggesting that the models' internal representations are highly tuned to the statistical properties of natural sequences.

These findings challenge the standard "last-layer" convention in computational biology, suggesting that researchers can achieve better results by systematically searching for the optimal layer rather than defaulting to the final one. The study provides a practical framework for feature selection in protein engineering and highlights that PLMs are not universal feature extractors; their utility is constrained by the alignment between their pre-training objectives and the specific downstream task at hand.