Structural Interpretations of Protein Language Model Representations via Differentiable Graph Partitioning

Authors: Siddhant Dutta, Edward Tan Beng Wai, Soumick Sarker et al.

arXiv · PDF

Summary

The authors build a framework to make protein language model (ESM-2) representations interpretable by projecting them onto protein contact graphs — networks where nodes are amino acid residues and edges represent spatial proximity. They use SoftBlobGIN, a graph neural network with differentiable clustering, to group residues into functional substructures and perform structure-aware prediction. This yields both strong performance (92.8% accuracy on enzyme classification, AUROC 0.983 on binding-site detection) and directly auditable explanations: the model recovers known active-site residues and catalytic contact patterns without any supervision on those sites.

Main takeaways:

• Combines language model embeddings (ESM-2) with graph neural networks over protein contact graphs to get structure-aware predictions.
• Differentiable "blob" clustering automatically groups residues into functional substructures; blobs containing active sites show 1.85× higher importance scores.
• GNNExplainer recovers biologically meaningful active-site residues and catalytic patterns post hoc.
• Improves binding-site detection from 0.885 AUROC (ESM-2 alone) to 0.983, showing structural explanations aren't recoverable from language-model features alone.
• Plug-and-play design: no retraining of the language model, adds only ~1.1M parameters.

Relevance

Not directly related to my persona/midtraining work — this is about protein structure and function prediction, though the general idea of projecting dense latent representations onto interpretable structures (their contact graphs vs. my activation steering/prompt space) has a loose conceptual parallel.

Abstract

arXiv:2605.10985v1 Announce Type: new Abstract: Protein language models such as ESM-2 learn rich residue representations that achieve strong performance on protein function prediction, but their features remain difficult to interpret as structural $&$ evolutionary signals are encoded in dense latent spaces. We propose a plug-$&$-play framework that projects ESM-2 representations onto protein contact graphs $&$ applies $\textbf{SoftBlobGIN}$, a lightweight Graph Isomorphism Network with differentiable Gumbel-softmax substructure pooling, to perform structure-aware message passing $&$ learn coarse functional substructures for downstream prediction tasks. Across enzyme classification, SoftBlobGIN achieves 92.8% accuracy $&$ 0.898 macro-F1. Unlike post hoc analysis of protein language models alone, our method produces directly auditable structural explanations: GNNExplainer recovers biologically meaningful active-site residues, spatially localized functional clusters, $&$ catalytic contact patterns. On binding-site detection, SoftBlobGIN improves residue AUROC from $0.885$ using an ESM-2 linear probe to $0.983$, indicating that these structural explanations are not recoverable from language-model features alone. Learned blob partitions provide an additional layer of interpretability by automatically grouping residues into functional substructures, with blobs containing annotated active-site residues showing $1.85\times$ higher importance than other blobs ($\rho{=}0.339$, $p{=}0.009$), without any active-site supervision. Our framework requires no retraining of the language model, adds only $\sim$1.1M parameters, $&$ generalises across ProteinShake tasks, achieving $F_{\max}$ of $0.733$ on Gene Ontology prediction $&$ AUROC of $0.969$ on binding-site detection. We position this as an interpretable structural companion to protein language models that makes their predictions more transparent $&$ auditable.