Source-linked AI summary
Interpreting Protein Language Model Embeddings via Orthogonal Projection for Protein Fitness Prediction
Paulo Yanez Sarmiento, Pia Francesca Rissom, Manuel Pfeuffer, Marco Simnacher, Jordan F. Safer, Sumaiya Iqbal, Henrike O. Heyne, Nadja Klein, Bernhard Y. Renard
TL;DR
PLM embeddings perform well for protein fitness prediction but are not directly interpretable, leaving unclear which biochemical properties they encode. The paper removes linear, higher-order, and interaction effects of tabular biochemical features through orthogonal projection, then evaluates downstream performance and explained variance. Removing these effects reduces classifier performance, while the features explain up to 29% of predicted-logit variance, supporting their relevance to the embeddings.
Problem
PLM embeddings achieve strong protein-fitness prediction performance but are not directly interpretable, leaving unclear which biochemical features they encode.
Method
The paper selects biochemical features and removes their linear, higher-order, and interaction effects from PLM embeddings using post-hoc orthogonalization, then compares classifiers and regresses predicted logits on those features.
Results
Removing biochemical-feature effects reduced average classifier performance across all four protein functions and PLMs, while the features explained 0.15 to 0.29 adjusted R2 of embedding-based predicted logits.
Takeaways & Limitations
The findings show that PLM embeddings encode patterns correlated with biochemical properties relevant to protein-fitness prediction, including higher-order and interaction effects.
Takeaways & Limitations
The analysis focuses on human-proteome single-mutation data and does not consider multiple mutations or certain assays lacking the UniProt reference sequence.
Abstract
from arXiv · showhide
Recently, there has been a growing adoption of protein language models (PLMs) in biomedical science. Their embeddings provide a rich numerical representation of protein sequences which achieve state-of-the-art performance on several downstream tasks including protein fitness prediction. However, PLM embeddings are not directly interpretable and, thereby, it remains unclear what features they encode. To gain insight into which biochemical properties of the protein are driving the prediction, we leverage an orthogonal projection technique that removes linear effects of known tabular features from embeddings and extend it to high-order and interaction effects. In this way, we remove the effects of interpretable biochemical features from PLM embeddings. In an ablation study, we show that this leads to a decrease in performance for a downstream classifier trained only on the embeddings to predict protein fitness. In an additional evaluation, we find that these biochemical features explain a substantial part of the variance in the predictions of this classifier. Hence, we can show that PLM embeddings encode patterns correlated with biochemical properties and quantify their contribution to predicting protein fitness. This computationally efficient approach is not limited to the features or embeddings considered here and is readily transferable to problem settings beyond protein fitness prediction.
1 Introduction
Protein fitness prediction matters because mutations can alter protein function, while PLM embeddings achieve strong downstream performance but remain difficult to interpret. The paper addresses this by removing biochemical-feature effects from embeddings with orthogonal projection, including higher-order and interaction effects.
- Motivation: Mutations can change protein fitness, making their functional impact important to estimate, especially because many clinical variants remain of unknown significance.Experimental measurement is expensive and time-consuming, motivating computational prediction methods.
- Motivation: PLMs provide rich sequence representations and have achieved state-of-the-art performance across downstream tasks.These representations are numerical vectors derived only from one-dimensional protein sequences.
- Research gap: PLM embeddings are difficult to interpret because it remains unclear which encoded features drive protein-fitness predictions.
- Approach: The paper applies post-hoc orthogonalization to remove biochemical-feature effects from embeddings and extends it beyond linear effects to higher-order and interaction effects.The approach is presented as computationally efficient for interpreting PLM embeddings.
2 Related Work
Prior work established that PLM embeddings capture biologically relevant information and support protein-fitness prediction, while orthogonal projection methods provide tools for separating tabular and latent effects. The paper builds on these strands to interpret PLM embeddings using biochemical features.
- PLM interpretation: Visual analyses, attention maps, and neighborhood methods have linked PLM embedding spaces to protein structure and biochemical properties.
- Protein fitness prediction: PLM-based protein-fitness methods include zero-shot scoring and supervised classifiers or regressors trained on embeddings.Some approaches augment embeddings with site-specific amino-acid features, while task-specific fine-tuning can improve embedding-based methods.
- Orthogonal projection: Semi-structured neural networks use orthogonal projection to address identifiability when combining tabular features with latent neural representations.
- Connection to this work: The paper applies these interpretability ideas to protein-fitness classifiers by comparing models trained on embeddings with and without projected tabular-feature effects.
- Orthogonal projection: Prior applications used orthogonal projection to remove confounding or metadata effects from neural embeddings and to separate higher- and lower-order model effects.
3 Method
The method embeds protein sequences, selects biochemical features, and projects their effects out of the embeddings before comparing downstream classifiers. It also controls for random information removal and regresses embedding-based logits on biochemical features to quantify explained variance.
- Data representation: Protein fitness is formulated as binary classification, with functioning proteins labeled 1 and nonfunctioning proteins labeled 0.Sequences are mapped by a PLM to d-dimensional embeddings, while corresponding biochemical measurements form the tabular feature matrix.
- Feature processing: L1-based feature selection identifies tabular biochemical features relevant to protein-fitness prediction before projection and downstream comparison.
- Projection: Orthogonal projection removes linear effects of selected tabular features from PLM embeddings, and a representation-based projection extends removal to higher-order and interaction effects.The projected embeddings are used to train a separate linear downstream classifier for comparison with the original embeddings.
- Ablation control: Randomized projections preserve marginal feature distributions while disrupting correlations, providing a control for information loss caused by dimensionality reduction.
- Variance analysis: An evaluation model regresses embedding-based predicted logits on selected tabular features or their representation and reports adjusted R2 on an unused hold-out test set.
- Reproducibility: The implementation and experiments are publicly available in a GitLab repository.
4 Experiments
The experiments evaluate protein fitness prediction on filtered ProteinGym subsets using tabular features, PLM embeddings, and wild-type-centered embedding differences. Models are assessed with cross-validation and controlled downstream classifiers to examine predictive performance and embedding structure.
- Data: 113k, 69k, 51k, and 18k mutated sequences remain for Organismal Fitness, Activity, Expression, and Binding, respectively, after filtering.The subsets contain 24, 15, 13, and 6 reference wild-types, respectively, and use stratified 5-fold cross-validation.
- Data: The tabular features describe substitution biochemistry and wild-type structure using angles, solvent accessibility, secondary structure, and plddt.Structural properties are derived from predicted wild-type structures in AlphaFoldDB.
- Representations: Embedding differences u − u0 represent changes from the reference wild-type and produce centered data that is more suitable for mitigating wild-type or assay bias.UMAP is used only for visualization; the centered differences are used as downstream-model inputs.
- Models: The analysis combines feature selection, an MLP representation of higher-order and interaction effects, and Transformer-based ESM embeddings.ESM-1v and ESM-2 models are mean-pooled at sequence level, with ESM-2 sizes of 35M, 150M, and 650M parameters.
- Downstream Task: A logistic regression with L2 regularization predicts fitness from tabular features, embeddings, concatenated inputs, or orthogonally projected embeddings.The downstream classifier is deliberately restricted in complexity to assess what the embeddings already encode; MCC evaluates performance and AUC selects regularization.
5 Results
PLM embeddings generally outperform tabular features for protein fitness prediction, while orthogonalizing biochemical-feature effects reduces embedding-based performance. These effects are supported by control experiments and variance analyses linking embeddings to biochemical properties.
- Performance Comparison: PLM embeddings achieved the highest performance across four protein functions, substantially exceeding tabular features or their representations.For ESM-1v 650M, MCC was 0.47 for Activity and 0.51 for Expression.
- Performance Comparison: 15 of 16 conditional-independence tests rejected no additional embedding information beyond tabular features after Bonferroni correction.The exception was ESM-2 35M embeddings for Activity.
- Ablation Study: Removing linear tabular-feature effects decreased average MCC for every protein function and PLM; for ESM-1v 650M, drops ranged from −0.09 to −0.14.The reported changes were −0.09 (0.39 →0.30), −0.13 (0.47→0.34), −0.14 (0.51→0.37), and −0.10 (0.41 →0.31).
- Ablation Study: Removing tabular-feature representations caused larger decreases, including an ESM-2 35M MCC decline from 0.35 to 0.15.The stronger effect indicates that higher-order and interaction effects relevant to protein fitness are encoded in the embeddings.
- Evaluation Model: Biochemical features explained 15% to 29% of the variance in ESM-1v 650M embedding-based predicted logits.Activity and Expression had the highest adjusted R2, while Binding had the lowest.
6 Conclusion
The study presents an efficient way to interpret PLM embeddings by removing biochemical features and testing their contribution to protein-fitness predictions. The results indicate that biochemical effects materially influence embedding-based predictions and that the approach can transfer beyond the analyzed features.
- Removing biochemical features from PLM embeddings substantially decreases the performance of classifiers predicting protein fitness.
- Random-subspace controls indicate that the performance decrease reflects removal of relevant biochemical features rather than random information or dimensionality reduction.
- Biochemical features, including higher-order and interaction effects, explain part of the variance in embedding-based classifier predictions.
- The computationally efficient approach is transferable to broader classes of biological features and problem settings.
Disclosure of Interests
The authors report no relevant competing interests and disclose using generative AI only for language editing and Python-code assistance.
- The authors declare no competing interests relevant to the article.
- Generative AI was used for language editing and Python-code development, not for data analysis, result interpretation, or scientific conclusions.
Appendix
The appendix provides figures and tables documenting predictive performance, orthogonal-projection comparisons, coefficient heatmaps, and regression-based evaluation results.
- Figure 6 compares MCC for linear classifiers trained on tabular features, their hidden-layer representations, PLM embeddings, and concatenated modalities.
- Figure 7 compares PLM embeddings with embeddings orthogonally projected against randomized or biochemical tabular features and their representations.
- Table 5 reports adjusted R2 for regressions of embedding-based predicted logits on tabular features or their representations across orthogonalization conditions.
- Table 6 reports average structural-feature coefficients across folds and identifies estimates significant across all folds.
- Figure 8 visualizes evaluation-model coefficients for single amino-acid substitutions and substitutions grouped by chemical property.