Intrinsic dataset features drive mutational effect prediction by protein language models

This study demonstrates that the predictive performance of protein language models on mutational effects is primarily driven by intrinsic dataset features, specifically site-level fitness variability, rather than model architecture, revealing that current models often fail to capture broader mutational patterns and frequently rely on simple site-mean baselines due to data leakage in standard benchmarks.

The Gist

The Big Picture: The "Magic" Protein Predictor That Isn't So Magic

Imagine you have a super-smart AI robot (called a Protein Language Model or pLM) that has read almost every book ever written about proteins. You ask it: "If I change this one letter in a protein's code, will the protein still work, or will it break?"

Scientists have been using this robot to predict how viruses and human cells react to mutations. Sometimes the robot is a genius; other times, it seems to be guessing randomly. This paper asks: Why is the robot so inconsistent?

The authors discovered that the robot isn't actually "thinking" deeply about the protein's chemistry. Instead, it's mostly just cheating by memorizing the location of the mutation. When the robot is tested on viruses, it fails because the viruses don't give it enough "cheat codes" to memorize.

---

The Core Problem: The "Address" vs. The "Message"

To understand the findings, let's use an analogy of a School Exam.

1. The "Pooled" Cheat (How we usually test)

Imagine a teacher gives a student a practice test. The test has questions about Math and History.

• The Cheat: The teacher mixes the questions up randomly. The student sees a question about "The Battle of Hastings" in the practice section and learns the answer. Then, on the final exam, the teacher asks about "The Battle of Hastings" again.
• The Result: The student gets a 100% score! But did they learn History? No, they just memorized the answer to that specific question.

In the paper, this is called Pooled Splitting. The AI sees mutations at "Site A" during training, and then gets tested on other mutations at "Site A." It doesn't learn how mutations work; it just learns that "Site A is usually bad" or "Site A is usually good." It memorizes the address, not the message.

2. The "Site-Stratified" Test (The honest way)

Now, imagine the teacher changes the rules.

• The Honest Rule: If the student sees "Site A" in the practice test, they are forbidden from seeing any questions about "Site A" on the final exam. The final exam only has questions about "Site B," "Site C," and "Site D"—places the student has never seen before.
• The Result: The student's score crashes. They realize they didn't actually learn the rules of History; they just memorized specific answers.

The paper shows that when scientists use this "Honest Rule" (splitting data by site), the AI's performance drops significantly. It proves the AI was mostly just memorizing site-specific averages, not learning the complex biology.

---

The Viral vs. Cellular Mystery

The authors noticed a weird pattern: The AI works okay on Cellular proteins (human/animal cells) but fails miserably on Viral proteins (like flu or HIV).

Why? They introduced two new "rulers" to measure the data:

Ruler 1: The "Variability of Addresses" (RVSM)

• Analogy: Imagine a city.
  • Cellular City: Some neighborhoods are very strict (always bad), some are very chill (always good), and some are chaotic. There is a big difference between neighborhoods.
  • Viral City: Almost every neighborhood is exactly the same. They are all "chill."
• The Finding: The AI loves the Cellular City. Because the neighborhoods are so different, the AI can easily guess, "Oh, this is a strict neighborhood, so this mutation is probably bad." It's an easy shortcut.
• The Viral Problem: In the Viral City, every neighborhood looks the same. The AI can't use its "neighborhood shortcut" because there are no distinct neighborhoods to memorize. It has to actually understand the chemistry, which it isn't very good at yet.

Ruler 2: The "Chaos Factor" (FHVS)

• Analogy: Imagine a classroom.
  • High Chaos: Every student behaves differently. Some are loud, some are quiet, some are sleepy.
  • Low Chaos: Everyone sits perfectly still.
• The Finding: The AI performs best when there is a Goldilocks zone of chaos.
  • If a site is too stable (Low Chaos), there's nothing to predict.
  • If a site is too chaotic (High Chaos), it's too noisy to learn patterns.
  • Viral proteins often have too many "stable" sites (Low Chaos). The mutations there don't change anything, so the AI has no signal to learn from.

---

The "Naive" Baseline: The Magic 8-Ball

The most surprising part of the paper is the "Naive Baseline."

The authors built a super-simple model that did nothing but look at the average score of a specific site.

• Example: "At position 50, mutations usually result in a score of 0.5. So, I predict 0.5 for any new mutation at position 50."

The Shock: On many viral datasets, this simple "Magic 8-Ball" performed just as well as, or better than, the super-complex, billion-parameter AI.

What this means: The complex AI wasn't doing anything special. It was just mimicking the simple average. The "intelligence" we thought the AI had was actually just the data itself telling the story.

---

The Takeaway: What Should We Do?

1. Stop Cheating: We need to stop testing AI models with "Pooled Splits" (where the AI sees the same location in training and testing). It gives us a false sense of security. We must use "Site-Stratified" splits to see if the AI can truly generalize.
2. Viral Proteins are Hard: Predicting mutations in viruses is much harder because they are evolutionarily "flexible" (mutations often don't matter). The AI struggles here not because it's broken, but because the data doesn't have clear patterns to learn.
3. The AI is Overhyped: For many tasks, the AI is just memorizing the "address" of the mutation. It hasn't truly learned the deep biochemical rules of life yet.

In short: The paper pulls back the curtain on the "magic" of protein AI. It tells us that the AI is often just a very good student who memorized the answer key, rather than a genius who understands the subject. To get real answers, we need to test it on questions it has never seen before.

Technical Summary

1. Problem Statement

Protein Language Models (pLMs) have become a standard tool for predicting protein fitness landscapes and mutational effects using Deep Mutational Scanning (DMS) data. However, their performance is highly inconsistent across different datasets.

• The Discrepancy: pLMs often perform significantly worse on viral datasets compared to cellular datasets, a trend that persists regardless of model architecture or transfer learning strategies.
• The Gap: The underlying causes of this performance variation are poorly understood. It is unclear whether the failure stems from model limitations (e.g., insufficient pretraining on viral sequences) or intrinsic properties of the datasets themselves.
• The Benchmarking Issue: Common evaluation practices (specifically "pooled" data splitting) may artificially inflate performance metrics by allowing data leakage, where the model learns site-specific averages rather than generalizable mutational rules.

2. Methodology

The authors conducted a systematic evaluation of supervised transfer learning across 74 DMS datasets (41 viral, 33 cellular).

• Models Evaluated:
  • Base Models: ESM-2 (650M parameters) and ESM C (600M parameters).
  • Domain-Adapted Models: Custom finetuned versions of ESM-2 on viral sequences (CRVDB and URVDB) and a publicly available viral-finetuned ESM-2 3B model (Sawhney).
  • External Benchmarks: Five top-performing models from the ProteinGym benchmark (including Kermut, ProteinNPT, MSA Transformer, Tranception, and ESM-1v).
• Transfer Learning Strategy:
  • Extracted mean-pooled embeddings from the final layer of pLMs.
  • Applied Lasso regression (and LoRA finetuning for specific experiments) to predict mutational fitness effects.
• Data Splitting Strategies:
  • Pooled Split: Mutations are randomly assigned to train/test sets, allowing the same sites to appear in both. This is the standard but potentially leaky approach.
  • Site-Stratified Split: All mutations at a specific site are assigned exclusively to either the training or test set. This forces the model to generalize to unseen sites, preventing the memorization of site-specific averages.
• Novel Metrics: The authors introduced two metrics to quantify dataset structure:
  1. Relative Variability of Site Means (RVSM): Measures the extent to which fitness variation is driven by differences between sites (site effects) versus variation within sites.
  2. Fraction of Highly Variable Sites (FHVS): The proportion of sites exhibiting significant within-site fitness variation (normalized standard deviation > 0.7).

3. Key Results

A. Viral vs. Cellular Performance Gap

• pLMs consistently underperform on viral datasets compared to cellular ones.
• Domain Adaptation Failure: Finetuning models on large viral sequence databases (RVDB) or using larger models (ESM-2 3B) did not close the performance gap. While it slightly improved viral performance, it often degraded cellular performance, and the difference remained statistically significant.
• ESM C Limitations: The newer ESM C model, which excels on cellular data, performed particularly poorly on viral data, likely due to the exclusion of viral sequences during its pretraining (safety concerns).

B. The Dominance of Site Effects (The "Naive" Baseline)

• A simple baseline model that predicts the fitness of a mutation based solely on the mean fitness of that site (calculated from training data) often matches or outperforms sophisticated pLMs, especially on viral datasets.
• This suggests that pLMs are frequently "cheating" by learning site-specific averages rather than capturing complex sequence-function relationships.

C. Impact of Data Splitting

• Pooled Splits: Inflate performance because the model sees the same sites in training and testing, effectively memorizing site means.
• Site-Stratified Splits: Cause a substantial drop in performance (often reducing $R^2$ to near zero) for all models. This confirms that models fail to generalize to unseen sites and rely heavily on site identity.
• Under site-stratified splits, the performance gap between viral and cellular datasets disappears, indicating the gap was an artifact of data structure and splitting methods, not inherent model inability to learn viral biology.

D. Predictive Power of Dataset Metrics

• RVSM: There is a strong positive correlation between RVSM and model performance. High RVSM (large differences between site means) makes datasets easier to predict because the model can rely on site identity.
• FHVS: Performance peaks at intermediate FHVS values.
  • Low FHVS: Most sites are invariant; mutations have no effect, making prediction trivial but uninformative.
  • High FHVS: High within-site variability but low between-site variability makes it hard to distinguish sites.
• Viral vs. Cellular Structure: Viral datasets typically have low FHVS (few sites vary significantly) and high RVSM, whereas cellular datasets have higher FHVS. This structural difference explains why viral datasets are harder for models to generalize on when site information is removed.

E. ProteinGym Validation

• Analysis of ProteinGym benchmarks confirmed that pooled splits yield significantly higher Spearman correlation ($\rho$) than stratified splits.
• The authors' RVSM and FHVS metrics successfully predicted model performance across ProteinGym datasets, explaining up to 61% of the variance in performance for the top model (Kermut).

4. Key Contributions

1. Identification of Dataset Drivers: Demonstrated that intrinsic dataset features (RVSM and FHVS), rather than model architecture or pretraining data volume, are the primary drivers of predictive success in mutational effect prediction.
2. Critique of Benchmarks: Provided strong evidence that the widely used pooled split strategy leads to data leakage, inflating performance estimates and masking the inability of models to generalize to unseen sites.
3. Baseline Superiority: Showed that in many cases (especially viral), a naive predictor using site means outperforms state-of-the-art pLMs, challenging the assumption that pLMs capture deep biochemical constraints.
4. Viral Data Challenge: Clarified that the poor performance on viral data is due to a lack of within-site variability (low FHVS) in the experimental data, not necessarily a lack of viral sequences in model pretraining.

5. Significance and Implications

• Re-evaluation of Model Claims: The field may be overestimating the generalization capabilities of pLMs. High performance scores in current literature often reflect the memorization of site-specific constraints rather than the learning of substitution-specific effects.
• Benchmarking Standards: The authors argue that site-stratified splits should replace pooled splits for benchmarking to ensure models are evaluated on their ability to generalize to new positions, which is critical for protein engineering applications.
• Experimental Design: Future DMS experiments should aim for a balanced spectrum of mutational effects (intermediate FHVS and RVSM) to generate data that is truly informative for machine learning.
• Model Development: Simply increasing model size or domain-adapting to viral sequences is insufficient. New architectures or input representations that better preserve position-specific information without relying on site identity memorization are needed.

Conclusion: The paper concludes that the "predictive signal" attributed to pLMs is largely an artifact of dataset composition and evaluation methodology. True progress requires rigorous evaluation strategies that prevent data leakage and acknowledge the limitations imposed by the variability structure of the underlying biological data.