Interpreting Latent Protein Language Model Features with Geometric Annotations

Siddharth Setlur; Djordje Mihajlovic; Darrick Lee

q-bio.QMAug 26, 2026

Interpreting Latent Protein Language Model Features with Geometric Annotations

Authors: Siddharth Setlur, Djordje Mihajlovic, Darrick Lee

Organizations: School of Mathematics University of Edinburgh, UK

Abstract

Protein language models (pLMs) encode information about protein sequences which enable downstream tasks such as structure prediction, but their internal representations are not well understood. Sparse autoencoders (SAEs) provide a promising tool to disentangle latent pLM representations into interpretable features, but existing annotation pipelines largely rely on protein-level annotations derived from database labels and LLM annotations of top activating sequences. Such annotations can overlook the localized residue-level and geometric patterns encoded by sparse features. We introduce an automated and scalable method for interpreting SAE features in ESM-2 by using geometrically inspired features of the protein $\text{C}_α$ backbone. Across ESM-2 8M layers, an FDR-controlled discovery analysis shows that local geometry is significantly associated with many SAE features, with varying levels of predictive strength, expanding coverage beyond database and sequence-based methods. In particular, geometry can distinguish SAE features sharing the same database annotation, revealing substructure within known biological labels. A significant portion of SAE features activate on unannotated metagenomic protein sequences enabling us to use our SAE annotations to better understand these sequences. In addition, ablation experiments at the level of contact prediction show that removing found geometric features shifts ESM-2's predicted contact maps in the direction of the descriptor. This provides a robust method of annotating proteins activated within SAE neurons at a residue level, providing a bridge between mechanistic interpretability and structural biology.

Figures & tables

Figure 1: Sparse autoencoder trained on the residual stream of ESM2-8M. The SAE encoder maps the 320-dimensional residual into 10,240 sparse features (

32\times

overcomplete), and the decoder reconstructs the original activation.

	% Total annotated	% InterPro Prot.	% InterPro Res.	% Seq Pos	% Seq Motif	% Geometric
Layer 2	86.53	56.11	42.03	62.22	69.33	84.38
Layer 4	93.50	56.75	61.46	79.19	77.56	92.79
Layer 6	94.57	56.68	67.70	71.35	77.20	94.49

Table 1: FDR-controlled discovery coverage across annotation methods. We say that a feature is annotated by a given annotation method, if its

q

value if below

0.05

. We say that a feature is annotated if any of the annotation methods have a

q

value if below

0.05

	PR-AUC 0.0-0.3	PR-AUC 0.3-0.6	PR-AUC $>0.6$
Layer 2	95.15	3.27	1.58
Layer 4	81.72	12.24	6.05
Layer 6	91.86	6.51	1.63

Table 2: PR-AUC values of geometrically annotated features on a set of 10k held out proteins across tested layers within ESM-2 8M. The 10k held out set is similarly sampled from the SwissProt database, explicitly disjoint from the 50k training samples.

Figure 2: Examples of different annotation methods across features. Comparing the annotations returned against all methods and respective statistical null tests (as in Table 1 ) we present SAE features whose annotation is considered geometry-primary (left) and InterPro residue primary (right). It is clear from the select examples and activations at their top activating proteins (plots accessed via [ 44 ] ) that the dominant annotation type can vary. We place specific emphasis on the feature f/894 whose annotation is only accurately available within the geometric pipeline developed in this work. In the case of the geometry-primary feature, the high GBM activation prediction accuracy automatically helps demonstrate that the underlying notion encoded at this feature is likely structural in nature.

	No. InterPro Res families	% Geom Distinguishable
Layer 2	120	59.17
Layer 4	578	77.51
Layer 6	45	68.89

Table 3: Geometric annotations further split InterPro residue annotations, revealing more granular features encoded across a subset of SAE features between layers.

Figure 3: Annotation exists beyond coarse InterPro labels. Here, we compare a set of SAE features whose biological annotation via residue-level InterPro label is consistently ‘Histidine kinase/HSP90-like ATPase domain’. Activation can be seen to vary substantially across the domain, suggesting finer features beyond the coarse biological annotation, which the geometric annotation method is able to detect. On this specific example, activation occurs across four ‘sub-domains’ within the larger IPR003594 annotation which remain distinct in both geometry space and sequence localization (median cosine 0.25; 80.6% zero overlap on shared proteins). Note, all GBM geometric feature vector predictions for these activations achieved PR-AUC $>0.65$ .

	% NMPFam act.	% geom. $q$ -sig.	% feat. med. PR-AUC $>$ 0.5	NMPFams matched	Sequences annotated
Layer 2	63.55	91.72	1.65% (169)	4.74% (2,369)	4.40% (440,475)
Layer 4	77.78	93.50	3.67% (376)	7.75% (3,875)	7.58% (757,802)
Layer 6	90.95	96.18	4.10% (420)	16.22% (8,108)	15.88% (1,588,446)

Table 4: Geometric annotation of metagenomic proteins. Col. 1: % of features with NMPFam activation. Col. 2: % of NMPFam activated features that are geometrically

q

-significant. Col. 3: % of those features with median PR-AUC

>0.5

evaluated on activating metagenomic sequences. Col. 4 % of 50k families with PR-AUC

>0.5

. Col. 5: % of 10M proteins with PR-AUC

>0.5

Figure 4: Local geometric annotations transfer to metagenomic proteins. Geometry predicts activation of a sample SwissProt protein $P37016$ well (left) and comparable perfomance on unseen metagenomic protein $F011972$ (right) with similar twisted beta sheet motifs consistently highlighted.

Figure 5: Contact prediction ablation experiments at layer 4. (Left) Ablation strength versus normalized signed target delta for SAE feature $f/\texttt{373}$ , revealing the linear relationship between response in dominant geometric feature vector features and ablation strength suggesting a causal relationship. (Middle) Contact probability at binned sequence separations versus ablation strength, revealing a clear, increasing relationship between probability of contacts at sequence separations [25-48] with increasing ablation strength averaged across 20 proteins firing at $f/\texttt{373}$ . (Right) Ablation experiments across 947 proteins covering the top 200 SAE features; a significant proportion see target geometric feature shift beyond matched-control tests.

Appendix figures & tables8 assets

Supplementary material from the paper’s appendix.

Appendix

Hyperparameter	Distribution	Range / values	Notes
Learning rate $\eta$	log-uniform	$[5\!\times\!10^{-5},\,10^{-3}]$	Adam initial LR
L1 coefficient $\lambda$	log-uniform	$[10^{-2},\,0.2]$	full strength after warmup
Dictionary size $m$	categorical	$\{2{,}560,\,5{,}120,\,10{,}240,\,20{,}480\}$	8 $\times$ , 16 $\times$ , 32 $\times$ , 64 $\times$
Batch size	categorical	$\{1{,}024,\,2{,}048,\,4{,}096\}$	residue tokens / step
Warmup ratio	uniform	$[0.05,\,0.15]$	of 500 000 total steps
Decay-start ratio	uniform	$[0.70,\,0.90]$	of 500 000 total steps

Appendix

Table 5: Hyperparameter sweep configuration.

Architecture
	Layer 2	Layer 4	Layer 6
Activation dim $d$	320	320	320
Dictionary size $m$	10 240	10 240	10 240
Expansion factor	32 $\times$	32 $\times$	32 $\times$
Optimisation
Batch size	4 096	2 048	1 024

Appendix

Table 6: Selected SAE per layer. All three SAEs are L1-ReLU autoencoders with the architecture of [ 34 ] and the training recipe of [ 17 ] .

Hyperparameter	Value
Number of estimators	80
Maximum tree depth	3
Learning rate	0.1
Subsample ratio	0.8
Minimum samples per leaf	$\max(5,\lfloor 0.02\times N\rfloor)$

Appendix

Table 7: Hyperparameter configuration for the Gradient Boosting Classifier.

Predicate name	Position condition	Description
first_5	pos < 5	First 5 residues (0–4).
first_10	pos < 10	First 10 residues (0–9).
first_20	pos < 20	First 20 residues (0–19); typical signal-peptide window.
last_5	pos >= slen - 5	Last 5 residues.
last_10	pos >= slen - 10	Last 10 residues.
last_20	pos >= slen - 20	Last 20 residues.

Appendix

Table 8: Details of the various position predicates that capture whether an SAE feature activates on a a certain position along the protein sequence. Note that

\mathtt{pos}

is the

0

indexed residue position, while

\mathtt{slen}

is the length of the protein sequence, so the conditions describe the positions that a particular predicate evaluates to true.

Figure 6: Top descriptor counts at layer 4. Top descriptor counts (GBM importance above 0.1) counted over all geometric annotations with PR-AUC

>0.3

showing vast majority of geometric annotations exhibit minimal dependance on amino acid covariates.

	% Total annotated	% InterPro Prot.	% InterPro Res.	% Seq Pos	% Seq Motif	% Geometric
Layer 6	92.99	65.68	64.28	81.52	70.08	92.92

Appendix

Table 9: FDR-controlled discovery coverage across annotation methods. We say that a feature is annotated by a given annotation method, if its

q

value if below

0.05

. We say that a feature is annotated if any of the annotation methods have a

q

value if below

0.05

	PR-AUC 0.0-0.3	PR-AUC 0.3-0.6	PR-AUC $>0.6$
Layer 6	63.20	26.30	10.50

Appendix

Table 10: PR-AUC values of geometrically annotated features across layer 6 within ESM-2 35M.

Protein	PDB	RMSD (Å)	Feature ID	nPos	PR-AUC (AF)	PR-AUC (crystal)
P75804	2g8s	0.13	1498	182	0.974	0.965
Q46927	4d79	0.43	6894	17	0.703	0.699
P0AC98	5zug	0.30	7257	16	0.725	0.694
Q47208	3dmy	0.44	2231	50	0.492	0.500
P0ADR8	9hs1	1.51	2231	28	0.488	0.516
P0ABE5	5oc0	0.44	6603	39	0.987	0.960

Appendix

Table 11: Comparison of PR-AUC values from layer 4 of ESM2-8M using AlphaFold (AF) and experimentally determined crystal structures. RMSD quantifies the structural difference between the AlphaFold model and corresponding PDB structure. The nPos column refers to the number of active residues for the protein-feature pair.

Mechanistic Interpretability

Mechanistic interpretability studies how neural networks implement behaviors through internal features, circuits, and causal information flow.

See all 391 papers

Structural Interpretations of Protein Language Model Representations via Differentiable Graph Partitioning
May 9, 2026Siddhant Dutta, Edward Tan Beng Wai, Soumick Sarker +2Graph Neural Networks Protein Language Models
Nanyang Technological University, Singapore
Structure, Not Sequence: Tracking and Validating Folding Mechanisms in ESMFold's Trunk with Sparse Autoencoders
Oct 26, 2026Bohan Liu, Xiaohan Fan, Wenqian Ye +1Protein Structure Prediction Neural Network Interpretability
University of Virginia Charlottesville VA 22903, USA
Interpreting Protein Language Model Embeddings via Orthogonal Projection for Protein Fitness Prediction
Aug 26, 2026Paulo Yanez Sarmiento, Pia Francesca Rissom, Manuel Pfeuffer +6Protein Language Models Protein Representation Learning
Hasso Plattner Institute, University of Potsdam · Bioinformatics and Machine Learning, Broad Institute of MIT and Harvard · School of Business and Economics, Humboldt University of Berlin +1

May 9, 2026J/K move · Enter open · S save

Structural Interpretations of Protein Language Model Representations via Differentiable Graph Partitioning

Siddhant Dutta, Edward Tan Beng Wai, Soumick Sarker, Pasan Gunawardane+1

Nanyang Technological University, Singapore

Graph Neural Networks Protein Language Models Protein Function Prediction Molecular Property Prediction Protein Representation Learning GNN Explainability

Open paper AskPDF 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

\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 (

ρ{=}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}

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.

Abstract

Mechanistic Interpretability

Explore similar work

Structural Interpretations of Protein Language Model Representations via Differentiable Graph Partitioning

Structure, Not Sequence: Tracking and Validating Folding Mechanisms in ESMFold's Trunk with Sparse Autoencoders

Interpreting Protein Language Model Embeddings via Orthogonal Projection for Protein Fitness Prediction