# Structural Interpretations of Protein Language Model Representations via Differentiable Graph Partitioning

arXiv: https://arxiv.org/abs/2605.10985 · HTML: https://mlsanity.com/paper/2605.10985/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 12 of its 20 references that are in the map.
- First posted on arXiv 9 May 2026.

## Builds on (most cited first, 12 of 12)

- [Semi-Supervised Classification with Graph Convolutional Networks](https://mlsanity.com/paper/1609.02907/) (2016)
- [Graph Attention Networks](https://mlsanity.com/paper/1710.10903/) (2017)
- [How Powerful are Graph Neural Networks?](https://mlsanity.com/paper/1810.00826/) (2018)
- [Axiomatic Attribution for Deep Networks](https://mlsanity.com/paper/1703.01365/) (2017)
- [Categorical Reparameterization with Gumbel-Softmax](https://mlsanity.com/paper/1611.01144/) (2016)
- [Strategies for Pre-training Graph Neural Networks](https://mlsanity.com/paper/1905.12265/) (2019)
- [Representation Learning on Graphs with Jumping Knowledge Networks](https://mlsanity.com/paper/1806.03536/) (2018)
- [Hierarchical Graph Representation Learning with Differentiable Pooling](https://mlsanity.com/paper/1806.08804/) (2018)
- [GNNExplainer: Generating Explanations for Graph Neural Networks](https://mlsanity.com/paper/1903.03894/) (2019)
- [Explainability in Graph Neural Networks: A Taxonomic Survey](https://mlsanity.com/paper/2012.15445/) (2020)
- [Protein Representation Learning by Geometric Structure Pretraining](https://mlsanity.com/paper/2203.06125/) (2022)
- [Learning Hierarchical Protein Representations via Complete 3D Graph Networks](https://mlsanity.com/paper/2207.12600/) (2022)

## Abstract (arXiv)

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 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× higher importance than other blobs (ρ=0.339, p=0.009), without any active-site supervision. Our framework requires no retraining of the language model, adds only ~1.1M parameters, & generalises across ProteinShake tasks, achieving Fmax 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.