Distilling Structural Representations into Protein Sequence Models

We have incorporated the suggested experiments, baselines and discussion into our revised paper draft (in red).

Incorporating MutRank in ISM

During training, ISM incorporates MutRank (MR) using structure-tuning. MutRank takes the atomic microenvironment as input and produces a single microenvironment encoding. The encoding is then clustered into one of K=512 token IDs, which then supervises ISM training. More specifically, we attach a linear MutRank decoding head to ISM and apply cross entropy loss to predict the MutRank token ID. In total, we train ISM with 3 decoding heads to predict the wild-type amino acid, Atomic Autoencoder token ID, and MutRank token ID.

During inference, MutRank is NOT run. Only the ISM sequence model is executed.

Comparisons with MutRank

We compare against (1) an ISM variant trained with only MutRank tokens (no Atomic Autoencoder) and (2) another ISM variant trained with two independent MutRank models (ISM MR^2). We obtain the second MutRank token IDs by retraining an independent MutRank.

On the structural benchmarks presented in Table 1, the ISM variant using only MutRank performs comparably to our final ISM model. However, on mutation effect prediction tasks in Table 2, while ISM with only MutRank shows improved performance, the full ISM model, incorporating both MutRank and the Atomic Autoencoder, outperforms it. Critically, only by incorporating the Atomic Autoencoder were we able to achieve performance on par with state-of-the-art methods which use multiple sequence alignments and structures as input.

We observe that ensembling two independent MutRank models does not bridge the performance gap in downstream mutation effect prediction. However, while this approach does not impact performance on structure or secondary structure prediction, it does enhance performance in contact prediction and binding residue tasks.

Comparisons with SaProt

While the difference in performance compared to SaProt may seem fairly marginal, SaProt explicitly takes the structure as input for downstream tasks. This requires users to have access to the structures and a sophisticated structure-based data pipeline. On the other hand, ISM requires only a sequence as input and is compatible with existing ESM software. We believe this makes ISM a valuable resource for the community despite its similar performance.

We additionally show that our model outperforms SaProt on mutation stability effect prediction with a Spearman correlation of 0.53 vs 0.49.

This review poorly summarizes the contributions of ISM and unfairly scores the paper low. While there are "structure-aware" pLMs, DistilProtBert is not one of them—this comparison is not justified. DistilProtBert is simply a distilled version of ProtBert and does not integrate structural tokens or use structural pretraining tasks. In contrast, ISM explicitly integrates structure tokens derived from microenvironments into the sequence model latent space—without requiring the input of actual PDB structures. ISM thus is novel for leveraging structural information implicitly while maintaining the the inherent usability of sequence-only models like ESM-2.

The reviewer's claim that the pretraining task leads to "data leakage" is factually incorrect. The Atomic Autoencoder generates structure tokens by clustering residue embeddings from a structural representation, but the pretraining itself is unsupervised. The structure tokens are derived solely from residues’ local atomic environments. If this constitutes data leakage, then all protein structural models, including AlphaFold and SaProt, would also be guilty of it. Furthermore, ISM operates sequence-only at inference time, meaning there is no reliance on structural inputs beyond the derived tokens used during pretraining.

The reviewer mentions "minimal performance improvement" without acknowledging the typical trend we see in protein modeling: performance improvements on top of high-performing models like ESM-2 are inherently incremental. ISM's improvements, while modest, translate into meaningful downstream benefits in structure-sensitive tasks like binding residue prediction and mutation stability assessment, where ISM provides consistently better results.

The use of multiple loss functions is a feature of ISM. Integrating structure prediction with sequence modeling is a well-established practice in multi-objective optimization of proteins. The manuscript extensively discusses how each component (e.g., structure tokens, structure-tuning loss) contributes to the overall training pipeline. The reviewer’s comment suggests they either overlooked or missed these explanations.

I will say this: ISM is novel in its approach to integrating structure tokens into sequence embeddings without relying on explicit structural inputs during inference. As a user and a developer of pLMs, I can say confidently that this makes ISM significantly more practical than models like SaProt, which require external structural data (e.g., PDB or AlphaFold predictions) at both training and inference stages.

Again, I want to be clear that ISM does introduce a well-motivated, novel strategy for structure-aware protein language modeling. The reviewer’s claims about data leakage and unmotivated design choices are unfounded. The scores provided are wholly unjustified and I encourage the Area Chair to consider this in their final decision.

Take GPT as an example: its task is defined as $x_{t+1} = f(x_1, x_2, ..., x_t)$, where the target $x_{t+1}$ is completely absent from the input.

For BERT, the task is $x_k = f(x_1, x_2, ..., x_{k-1}, \text{[MASK]}, x_{k+1}, ...)$, where the target $x_k$ is also unknown in the input. Although BERT typically applies masking to 15% of the tokens, I’ve simplified it to just one token here for clarity.

In contrast, for the "Atomic Autoencoder," the task is $x_{\text{coord}} = f(d_{11}, d_{12}, ...)$, which is a trivial task because the encoder already has access to the full structural information (from pair-wise distance $d$), making it straightforward to recover the coordinates $x_{\text{coord}}$. (As discussed in the above thread)

Furthermore, I don't think the "Atomic Autoencoder" fits the definition of a classical autoencoder. Classical autoencoders typically include an information bottleneck, such as dimensionality reduction, to make the learning process more challenging and meaningful. However, the "Atomic Autoencoder" lacks this feature, functioning more like a straightforward encoder-decoder system.

Performance improvements are limited: ISM vs. MutateEverything (AF)

Performance comparisons between ISM and MutateEverything (AF) are nuanced: ISM underperforms MutateEverything (AF) on Spearman correlation (0.53 vs. 0.56) and RMSE (1.44 vs. 1.38) but outperforms on MCC (0.40 vs. 0.35).

Importantly, MutateEverything (AF) relies on a multiple sequence alignment (MSA) as input, which encodes explicit evolutionary information but is computationally expensive to generate and may be unavailable for certain proteins, such as de-novo designs. In contrast, ISM requires only the target sequence as input, offering greater accessibility and versatility. For example, on a 317 amino acid protein, ISM is 20x faster than the AlphaFold-based model (0.6 vs 12.1 seconds).

Note that while AlphaFold was trained on fewer protein sequences in PDB compared to UniRef, it is trained on much richer supervision in the form of experimental structures.

Performance improvements are limited: Functional Tasks in Table 3

Our experiments demonstrate ISM's versatility and effectiveness across different tasks. On structural benchmarks and mutation effect prediction (Tables 1 and 2), ISM achieves superior performance compared to existing models. For other benchmarks where structural representations may be less critical (Table 3), ISM performs competitively with state-of-the-art models. These comprehensive results establish ISM as a robust choice for a wide range of protein modeling tasks, even ones requiring less structural knowledge.

Unique capabilities or advantages ISM provides compared to existing pLMs. ISM’s unique capability lies in its structural-enriched representations from a single sequence input. It does so without computationally or experimentally generating a structure while matching performance of structure-based models on downstream tasks.

The potential impact on the community of protein studies is not clearly presented.

Integrating structural information into protein language models (pLMs) while maintaining sequence-only input requirements represents a significant advancement for computational biology. We demonstrate the practical impact of this approach on two structure-dependent applications: (1) improving ESMFold's structure prediction capabilities, and (2) enhancing mutation effect prediction for ΔΔG. This advancement is particularly valuable because many computational biology applications rely on readily available sequence data while being implicitly influenced by unavailable protein structures. ISM's ability to capture structural information from sequence alone makes it applicable to these scenarios.

Discussion and Comparisons with AlphaFold as structure learner.

While AlphaFold requires both protein sequence and multiple sequence alignment inputs, ISM learns structural representations directly from sequence alone. This fundamental difference makes AlphaFold an unequal comparison and not the most relevant baseline. Nevertheless, for ΔΔG mutation effect prediction in Table 2, ISM matches the performance of state-of-the-art models, including MutateEverything (AF), which leverages AlphaFold as the backbone.

Discussion and Comparisons with ESMFold as structure learner.

ESMFold is a structure predictor which takes a protein sequence as input, derives a latent representation from a frozen ESM-2, then trains a folding trunk on top of the ESM2 representations. We directly compare against an ESMFold variant using the 650M parameters ESM-2 model in Table 1 and find that our model outperforms ESMFold with GDT-TS of 0.67 vs 0.64. Additionally, we compare ISM against ESM-2 representations for all structural benchmarks and found consistent performance improvements in Table 1.