Shuibai12138/off-manifold-collapse-plm

Off-Manifold Collapse in Guided Protein Language Models

Reference activation statistics for Mahalanobis filtering, the training-free post-hoc selector introduced in the paper. This repository contains only the 10 KB natural-activation statistic needed to score a candidate sequence. No model weights are included.

Abstract

Protein language models are widely used priors for protein sequence design, and a growing body of work controls them at inference time as an alternative to fine-tuning. Such guidance faces a dilemma: mild enough to preserve natural activation statistics, it barely moves the property; strong enough to move it, the generations become progressively harder to fold. We show the failure has a specific and cheaply detectable signature, an off-manifold collapse of the model's own representations. Guided activations fall toward a region statistically indistinguishable from random amino-acid input, and the sequences degenerate to low complexity, yet the property oracle being optimized can still score these generations as a success. The optimized oracle can therefore fail to witness the collapse and, for solubility, can actively reward it, whereas structure and composition expose the failure. Because the failure is already visible in a finished candidate, we detect it at the output rather than modify the generator. We introduce a cheap density prior over natural protein activations and keep only the candidates that remain typical under it, a training-free post-hoc step we call Mahalanobis filtering. At matched guidance settings it improves both the property score and the structural plausibility of the sequences it keeps at negligible cost, without touching the generator, and transfers across different guidance methods.

Contents

File	Size	Description
rep_statistics.pt	12 KB	dict with mean and var, each a float32 tensor of shape (1, 1280)

• mean — per-channel mean $\mu_d$ of natural layer-17 activations
• var — per-channel variance $\sigma_d^2$ of the same activations

These two 1280-dimensional vectors are the method's only learned component.

How the statistics were computed

Setting	Value
Base model	ESM-2 650M (esm2_t33_650M_UR50D)
Layer	17 (of 33)
Hidden dimension	D = 1280
Corpus	UniRef50, ~58M sequences of length 30–1022
Tokens used	interior tokens only (BOS/EOS excluded)
Estimator	single streaming Welford pass (corpus never materialized or revisited)
Token dropout	disabled

A single forward pass per sequence extracts layer-17 hidden states; per-channel first and second moments are accumulated online with Welford's algorithm. The result is estimated once, offline and then reused unchanged across every task, property, steering setting, guidance mechanism, and seed reported in the paper.

How they are used

Given the interior-token layer-17 activations $h \in \mathbb{R}^{L \times D}$ of a finished candidate, the typicality score is a diagonal Mahalanobis distance:

$${\mathrm{M}\mathrm{a}\mathrm{h}\mathrm{a}\mathrm{l}}^2(h)=\frac{1}{L}\sum _{i=1}^L\sum _{d=1}^D{\left(\frac{h_{i,d}-{\mu }_d}{{\sigma }_d}\right)}^2\backslash mathrm\{Mahal\}^2(h)\; =\; \backslash frac\{1\}\{L\}\backslash sum\_\{i=1\}^\{L\}\backslash sum\_\{d=1\}^\{D\}\backslash left(\backslash frac\{h\_\{i,d\}-\backslash mu\_d\}\{\backslash sigma\_d\}\backslash right)^2$$

Under a diagonal-Gaussian working model, the per-token sum has reference distribution $\chi^2(D)$ with mean $D = 1280$, so natural activations concentrate near that value. Guided generations that have collapsed off-manifold fall far below it (toward the random-amino-acid level of ~600). The filter accepts a candidate iff

$${\mathrm{M}\mathrm{a}\mathrm{h}\mathrm{a}\mathrm{l}}^2(h)\ge D-k\sqrt{2D},k=1\Rightarrow \text{threshold }1229.\backslash mathrm\{Mahal\}^2(h)\; \backslash ;\backslash ge\backslash ;\; D\; -\; k\backslash sqrt\{2D\},\; \backslash qquad\; k\; =\; 1\; \backslash ;\backslash Rightarrow\backslash ;\; \backslash text\{threshold\; \}\; 1229\; .$$

The threshold is $\chi^2$-inspired, not $\chi^2$-calibrated: token positions are correlated, so the sequence-averaged score is not itself $\chi^2(D)$ and $k$ is an empirical margin on the natural scale. Results are insensitive to $k$ over $[0.3, 2]$.

Usage

import torch
from huggingface_hub import hf_hub_download

path = hf_hub_download("Shuibai12138/off-manifold-collapse-plm",
                       "rep_statistics.pt")
stats = torch.load(path, map_location="cpu")
mu    = stats["mean"].squeeze()          # (1280,)
sigma = stats["var"].squeeze().sqrt()    # (1280,)

def mahal2(h):
    """h: (L, 1280) interior-token layer-17 activations of one sequence."""
    z = (h.float() - mu) / sigma
    return (z ** 2).sum(-1).mean()

D, k = 1280, 1.0
threshold = D - k * (2 * D) ** 0.5       # 1229.4
accept = mahal2(h) >= threshold

Scoring costs one ESM-2 forward pass of the finished sequence (already available during generation) plus a 10 KB dot product: ~34 ms per sequence on an A100, compared with ~3 s for an exact flow-based density on the same activations.

Citation

@article{offmanifoldcollapse2026,
  title  = {Off-Manifold Collapse in Guided Protein Language Models},
  year   = {2026}
}