💻 computer science

Off-Manifold Collapse in Guided Protein Language Models

This paper identifies "off-manifold collapse," a failure mode in guided protein language models where strong guidance produces low-complexity, un-foldable sequences that fool property oracles, and proposes a training-free "Mahalanobis filtering" post-processing step to detect and remove these invalid candidates by filtering for natural activation statistics.

Original authors: Shuibai Zhang, Xinchi Liu, Fred Zhangzhi Peng, Zhihan Yang, Shutong Wu, Yingzi Ma, Jiawei Zhang

Published 2026-08-20

📖 6 min read🧠 Deep dive

CC BY 4.0

Original authors: Shuibai Zhang, Xinchi Liu, Fred Zhangzhi Peng, Zhihan Yang, Shutong Wu, Yingzi Ma, Jiawei Zhang

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Proteins are the molecular machines that build and run every living thing. They are long chains of building blocks called amino acids, and the specific order of these blocks determines how the chain folds into a three-dimensional shape. That shape is what allows a protein to function, whether it is digesting food, fighting an infection, or transmitting a signal. For decades, scientists have tried to design new proteins from scratch, but the space of possible sequences is so vast that finding a useful one is like searching for a needle in a haystack. In recent years, artificial intelligence has offered a new way forward. Large computer models, trained on millions of natural proteins, have learned the hidden rules of how these chains behave. These models can now generate new protein sequences that look and act like natural ones, serving as a powerful tool for creating medicines and industrial enzymes.

However, a new study reveals a critical flaw in how scientists currently try to tweak these AI models to create proteins with specific new traits, such as higher solubility or better heat resistance. The researchers found that when they push the AI too hard to achieve a desired property, the model stops producing realistic proteins. Instead, it collapses into a state where the generated sequences look like random noise to the model's own internal logic, even though a separate scoring system still rates them as successful. The team discovered a simple, fast way to spot this failure after the AI has finished its work, allowing them to filter out the broken designs without needing to retrain the model or change how it generates sequences.

The problem arises when researchers try to guide these language models toward a specific goal. Imagine a model that has read every known protein sequence and understands the subtle patterns that make them stable. Scientists can nudge this model during the generation process by adding a specific direction to its internal calculations, effectively telling it, "Make this protein more soluble." This technique, known as activation steering, is popular because it does not require the complex and expensive process of retraining the entire model. But there is a hidden cost. When the steering force is too strong, the model loses its grip on the natural patterns it learned. It begins to produce sequences that are statistically indistinguishable from a random string of amino acids.

The researchers observed this failure clearly when they tested the model on solubility. Under mild guidance, the AI produced proteins that folded well and had improved solubility. But as they increased the strength of the steering, the quality of the generated proteins plummeted. The model began emitting long chains made almost entirely of a single amino acid, glycine. These chains were so simple and repetitive that they could not fold into a functional shape. Yet, the computer program used to judge solubility still gave these broken chains a perfect score. The model had been tricked into thinking that a random, collapsed sequence was a success because the scoring system was looking for a specific signal that the random noise happened to mimic. The AI had moved off the "manifold," a term scientists use to describe the curved surface where all natural, functional proteins live, and fallen into a region of statistical randomness.

To understand why this happened, the team looked inside the model's own "brain" while it was generating these sequences. They examined the mathematical representations the model created for each amino acid. In a healthy generation, these representations stay within a specific range of values that the model learned from natural proteins. When the steering became too aggressive, these values shrank and collapsed toward the average, losing the unique variations that define a real protein. The researchers found that this collapse was a reliable warning sign. Even before the protein sequence was fully formed, the model's internal state was already showing signs of failure, becoming indistinguishable from the state it would have if it were just guessing random letters.

The solution the team proposed is surprisingly simple and requires no new training. They developed a filter that acts as a final check on any protein the AI generates. After the model finishes creating a sequence, this filter calculates a single score based on how typical the sequence's internal representations are compared to natural proteins. If the score is too low, indicating the sequence has collapsed into randomness, the filter discards it. If the score is high enough, the sequence is kept. This process is incredibly fast, taking only a fraction of a second per protein, and it uses a tiny amount of computer memory. It does not change how the AI generates the proteins; it simply selects the good ones from the pile of bad ones that the AI produced.

When the researchers applied this filter to their experiments, the results were striking. For the same level of steering force, the filtered set of proteins had much higher structural quality than the unfiltered set. The proteins were more likely to fold into stable shapes, and they still retained the improved properties the researchers were looking for. The filter worked just as well for different types of guidance, including methods that use gradients to nudge the model, proving that the problem was not unique to one technique. The team showed that by simply rejecting the sequences that had fallen off the natural path, they could recover the high-quality designs that the AI had accidentally buried under a pile of statistical noise.

This finding changes how scientists should approach protein design with AI. It suggests that the most obvious sign of success—a high score from a property predictor—is not enough. A sequence can score perfectly on a computer test while being structurally useless. The researchers demonstrated that the model's own internal state holds the truth about whether a design is real or fake. By listening to that internal state, they can separate the signal from the noise. The work does not claim to have solved the problem of protein design entirely, nor does it replace the need for real-world laboratory testing. Instead, it offers a practical, low-cost tool to ensure that the digital designs generated by AI are actually worth testing in the lab, preventing researchers from wasting time on sequences that are nothing more than mathematical illusions.

Technical Summary: Off-Manifold Collapse in Guided Protein Language Models

Problem Statement

Protein language models (PLMs) like ESM-2 serve as powerful priors for protein sequence design. A growing trend involves controlling these models at inference time via activation steering—adding a direction vector derived from high- and low-property reference sequences to the model's residual stream—rather than fine-tuning. However, this approach faces a critical dilemma: steering must be mild to preserve natural activation statistics, yet mild steering barely shifts the target property. Conversely, strong steering successfully moves the property score but causes generated sequences to degenerate.

The paper identifies a specific failure mode: Off-Manifold Collapse. When steering is pushed aggressively to maximize a property (e.g., solubility or thermostability), the model's internal representations collapse toward a region statistically indistinguishable from random amino acid input. This results in low-complexity sequences (e.g., glycine homopolymers) that fold poorly. Crucially, the property oracle used for steering often fails to detect this collapse; it may continue to score these degenerate sequences as highly successful, creating a false positive where the "optimized" sequence is structurally implausible.

Methodology

Diagnosis: Activation Collapse

The authors diagnose this failure by analyzing the generator's hidden states rather than just the output sequences. They define a diagonal Mahalanobis score ( $Mahal^2$ ) to measure how typical a sequence's activations are relative to natural proteins.

Reference Statistics: They estimate per-channel mean ( $\mu$ ) and variance ( $\sigma^2$ ) of activations from the UniRef50 dataset (natural proteins) at a specific layer (Layer 17 of ESM-2 650M).
The Metric: For a candidate sequence with hidden states $h$ , the score is calculated as the sum of squared standardized deviations:
$Mahal^2(h) = \frac{1}{L} \sum_{i=1}^{L} \sum_{d=1}^{D} \left( \frac{h_{i,d} - \mu_d}{\sigma_d} \right)^2$
Observation: As steering strength ( $\alpha$ ) increases, $Mahal^2$ decreases, indicating that activations contract toward the channel-wise mean. At high steering strengths, the score approaches that of random amino acid inputs, signaling a loss of "protein-like" structure in the representation space.

Solution: Mahalanobis Filtering

Instead of modifying the generator or the steering mechanism (which proved inconsistent across properties), the authors propose a training-free, post-hoc selection layer called Mahalanobis filtering.

Mechanism: After sequences are generated, the filter calculates the $Mahal^2$ score for each candidate.
Thresholding: A candidate is accepted if and only if:
$Mahal^2(h) \geq D - k\sqrt{2D}$
where $D$ is the embedding dimension (1280) and $k$ is a constant set to 1. This threshold is derived from the properties of a $\chi^2$ distribution but calibrated empirically against the natural distribution of unguided sequences.
Scope: This is a one-sided test targeting "inward collapse" (activations becoming too typical/random). It does not filter "outward" extremes.
Cost: The filter requires only a single forward pass of the completed sequence through the base PLM (ESM-2) and a dot product with stored statistics (10 KB), adding negligible computational cost (~34 ms per sequence).

Key Contributions

Characterization of Failure: The paper identifies and characterizes "off-manifold collapse" as the primary failure mode of aggressive activation steering. It demonstrates that property scores can increase while structural plausibility (measured by ESMFold pLDDT) and activation typicality ( $Mahal^2$ ) simultaneously degrade.
Mahalanobis-Filtered Steering: The authors introduce a simple, training-free post-hoc filter that selects only those generated candidates that remain "typical" under the natural activation distribution.
Transferability: The filter is shown to be effective across different guidance methods (activation steering and predictor-gradient guidance) and different target properties (solubility and thermostability) without retraining or recalibration.

Results

The evaluation was conducted on ESM-2 650M using UniRef50 references.

Trade-off Recovery: At matched steering strengths (e.g., $\alpha=0.5$ $α = 0.5$ ), applying the filter significantly improved structural confidence without sacrificing the achieved property shift.
- Solubility: Mean pLDDT increased from 0.401 to 0.647 (+0.247), while the predicted solubility probability also increased within the accepted subset.
- Thermostability: Mean pLDDT increased from 0.446 to 0.633 (+0.187), with a concurrent increase in predicted melting temperature ( $T_m$ ).
Filtering Efficiency: The filter retains a subset of the population (e.g., 16% for solubility at $\alpha=0.5$ ) but ensures that the retained sequences are structurally plausible.
Generalization: The filter improved pLDDT across all tested guidance strengths and mechanisms (including predictor-gradient guidance), confirming that the inward collapse is a general phenomenon of aggressive inference-time control.
Efficiency: The Mahalanobis statistic correlates well ( $r=0.82$ ) with a more expensive flow-based density model (GLP) but is orders of magnitude faster (34 ms vs. ~3 s) and requires no additional learned models.

Significance and Claims

The paper claims that the failure of guided PLMs is detectable in the finished candidate's activations rather than requiring changes to the generation process.

Detection over Correction: Because the "collapse" signature is visible in the output, the authors argue for a "detect and discard" strategy (filtering) rather than a "fix the generator" strategy. This is particularly valuable because generator-side repairs were found to be property-specific and difficult to tune.
Reliability of Oracles: The work highlights that property oracles can be "blind" to structural collapse, potentially rewarding degenerate sequences. The Mahalanobis filter acts as a necessary referee to ensure that property improvements are not bought at the cost of structural validity.
Modesty: The authors explicitly state that this method is a computational pre-screen. It selects from a viable tail of the generated pool but cannot manufacture candidates the generator did not produce. It does not claim to solve all failure modes (e.g., outward collapse) and relies on ESMFold pLDDT as a proxy for structure, not experimental validation.

In summary, the paper provides a cheap, generalizable diagnostic and mitigation strategy for the "off-manifold collapse" inherent in aggressive activation steering, ensuring that designed proteins remain structurally plausible while still achieving desired functional shifts.