Interpreting Latent Protein Language Model Features with Geometric Annotations

Siddharth Setlur; Djordje Mihajlovic; Darrick Lee

Source-linked AI summary

Interpreting Latent Protein Language Model Features with Geometric Annotations

Siddharth Setlur, Djordje Mihajlovic, Darrick Lee

arXiv:2608.26419v1Aug 26, 2026q-bio.QMcs.LG

TL;DR

Existing pLM interpretability methods can miss residue-localized and geometric patterns in SAE features. This paper introduces a scalable geometry-based annotation pipeline and shows refined geometric motifs, metagenomic transfer, and contact-map shifts under ablation.

Problem
pLM internal representations remain poorly understood, while protein-level annotations can miss residue-localized and geometric patterns in SAE features.
Method
The authors integrate residue-level database, sequence-motif, sequence-position, and geometrically informed annotations evaluated with a common significance-testing procedure.
Results
The pipeline expands annotation coverage, separates shared InterPro labels into geometric motifs, transfers some SwissProt annotations to metagenomic proteins, and shifts contact maps under ablation.
Takeaways & Limitations
Residue-level geometric annotations provide a framework for probing structure-related pLM representations and connecting mechanistic interpretability with structural biology.
Takeaways & Limitations
The analysis focuses on ESM-2 8M and predicted structures, so ablations are limited to contact-level interventions and may involve circularity from AlphaFold geometry.

Abstract

from arXiv · show

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.

1 Introduction

Protein language models learn biologically useful representations from sequence data, but their internal features remain difficult to interpret. This work introduces residue-level geometric annotations for sparse autoencoder features and shows that they expand coverage, refine labels, and transfer to metagenomic proteins.

Motivation: pLMs learn representations supporting structure and function prediction, yet the features they learn from unlabelled sequences and how they are used remain unclear.Mechanistic interpretability is presented as a way to identify human-interpretable features and support protein annotation and related applications.
Existing approaches: SAEs decompose dense model representations into larger sets of sparse features intended to be more interpretable than polysemantic individual neurons.Existing annotation approaches identify commonalities among strongly activating tokens or proteins, including database labels and LLM-generated descriptions.
Approach: The proposed pipeline tests whether SAE activations are predictable from local Cα backbone geometry, supplemented by amino-acid covariates and complementary database, motif, and positional annotations.The method operates at residue-level resolution and evaluates annotation signals with a common significance-testing framework.
Findings: Geometry annotates SAE features missed by InterPro, separates features with the same InterPro label into finer geometric motifs, and transfers annotations from SwissProt to some metagenomic proteins.These findings extend annotation beyond database labels alone and reveal substructure within known biological categories.
Findings: Geometry-annotated features are linked to contact-map changes under ablation, connecting residue-level SAE annotations with ESM-2 structural predictions.The introduction identifies contact-map changes under ablation as a specific contribution of the pipeline.

2 Background

ESM-2 maps protein sequences to contextual residue representations learned through masked-language modeling, while SAEs decompose these representations into sparse features. The study trains separate overcomplete ReLU SAEs across ESM-2 8M residue layers and evaluates their activations on diverse protein sequences.

Protein language models: ESM-2 is an encoder-only transformer that embeds each residue as a token and updates representations through bidirectional encoder layers.The model is trained on UniRef50 protein sequences.
Protein language models: Approximately 15% of residues are masked during training, and ESM-2 predicts the hidden amino acids from the remaining sequence context.Despite sequence-only training, its representations encode structural and functional information, including residue-contact information.
Sparse autoencoders: SAEs map an m-dimensional activation into an overcomplete F-dimensional feature representation and decode it back toward the original activation.The basic ReLU architecture is trained for faithful reconstruction while enforcing sparsity, producing SAE features intended to be monosemantic.
Experimental setup: Each of the six ESM-2 8M residue layers receives a separate ReLU SAE with 10,240 features, 32 times the 320-dimensional embedding.Features are learned independently at every layer and can activate at individual residue positions.
Experimental setup: SAE inference is performed on 50,000 protein sequences sampled from sequence-identity clusters so near-identical families do not dominate the analysis.Clustering uses 30% sequence identity and 80% bidirectional coverage, with at most one accession per cluster representative.

3 Residue-Level Annotation for pLM SAE Features

The paper presents a residue-level pipeline that tests whether SAE activations are predictable from local Cα backbone geometry alongside sequence, database, motif, and position annotations. It evaluates these predictions against active-residue labels using feature-wise scoring and FDR-controlled significance testing.

The pipeline evaluates geometric, database, sequence-motif, and sequence-position annotations using a common residue-level prediction framework.Each method produces a continuous score or binary predicate that is compared with SAE activations.
For each feature, the method selects proteins, records residue activations, thresholds active residues, and computes annotation predictions.Protein selection uses top-activating proteins for continuous methods or sampling across normalized activation levels for binary methods.
3.3 Evaluation and Significance Testing: Continuous methods use PR-AUC at a fixed threshold, whereas binary methods optimize F1 across thresholds before significance testing.The analysis compares scores with residue-shuffled null distributions, uses 100 permutations, and applies Benjamini–Hochberg correction across 10,240 features per method and layer.
3.1 Geometrically-Informed Annotation: Local geometric annotations use 21-residue windows and 44-dimensional vectors containing curvature, torsion, planarity, contact statistics, and amino-acid composition covariates.The windows are centered on each residue and extracted from AlphaFold DB structures.
3.1 Geometrically-Informed Annotation: Gradient-boosted classifiers predict active residue-level SAE activations from geometric feature vectors, with performance evaluated using PR-AUC.The geometric method tests whether residues on which an SAE feature activates are predictable from local descriptors rather than structural alignment.

4 Results

Across ESM-2 8M, residue-level geometry broadens SAE-feature annotation beyond database and sequence methods, identifies finer structure within shared InterPro labels, transfers to metagenomic proteins, and links geometric features to contact-map changes.

Annotation coverage: FDR-controlled testing found local geometry provided the broadest single-method coverage across SAE features, with annotation defined by permutation-test q-values below 0.05.Overall coverage was the union of database, sequence-motif, sequence-position, and geometric methods.
Annotation coverage: Geometric annotations capture localized structural patterns that can differ from broad InterPro residue annotations across SAE features.Representative geometry-primary and InterPro-primary features show that dominant annotation types vary, while some features receive overlapping annotations.
Validation: Similar results were obtained for ESM-2 35M layer 6, and a small subset of protein-feature pairs was additionally validated using experimentally determined structures.These evaluations are reported in Appendices H and I.
Residue-level substructure: Geometric feature vectors distinguish four SAE features sharing the same Histidine kinase/HSP90-like ATPase domain label by their distinct residue-localized subdomains and motifs.For the example in Figure 3, the features have median cosine similarity 0.25 and 80.6% zero overlap on shared proteins; all GBM predictions achieved PR-AUC > 0.65.
Contact prediction ablations: Ablating contact-related feature f/373 produced ablation-strength-dependent changes in its target contact-geometry metric and reduced predicted contacts at intermediate sequence separations.The strongest geometric predictor was mean sequence separation at 12Å, and the localized effect appeared in bins [13, 48].

5 Discussion

The discussion frames geometric SAE annotations as a way to support protein analysis beyond known labels and to probe structure-related representations. It also identifies scale, predicted-structure circularity, metagenomic confidence, and aggregate-validation limits.

Limitations: The analysis is limited to ESM-2 8M and contact-level interventions, rather than validation through a full folding head.The authors also note potential circularity from using AlphaFold-predicted geometry.
Limitations: Lower pLDDT makes geometric analysis of metagenomic proteins a less robust information source.
Downstream applications: Geometric SAE features may provide coarse motif-based similarity for novel sequences, helping biologists prioritize proteins for comparison and drug-development research.
Downstream applications: Geometric annotations may identify candidate seam positions for recombination by resolving InterPro domains at finer detail without requiring known three-dimensional coordinates.
Conclusion: The pipeline refines biological labels into local motifs, transfers annotations to some metagenomic proteins, and shifts contact maps under targeted feature ablation.

A SAE training details

The appendix describes SAE training and annotation-pipeline implementation, including data sampling, activation thresholding, computational resources, and model constraints.

SAE training: SAEs are trained for ESM2-8M residual-stream layers 2, 4, and 6 using an established architecture and training recipe.
SAE training: 50,000 uniformly sampled SwissProt sequences of length ≤1024 provide per-residue 320-dimensional activation vectors for SAE training and held-out fidelity evaluation.
SAE training: Decoder dictionary vectors remain unit ℓ2-norm through gradient projection and column renormalization after each optimizer step.
SAE training: The encoder and decoder use separate learned parameter matrices rather than tied weights.
Annotation pipeline: Continuous annotation methods sample top-activating proteins, whereas binary methods stratify proteins into one inactive and ten activation bins.
Annotation pipeline: Continuous methods define positive residues above each feature’s nonzero-activation 0.8 quantile, while binary methods sweep thresholds for F1 evaluation.

B.2 Details for Residue Level Database Methods

Residue-level database methods use activation-stratified protein samples and optimize per-label activation thresholds for InterPro-based evaluation.

Sampling strategy: Database-method sampling selects up to 50 proteins from each of 11 normalized-activation bins, for at most 550 proteins per feature.
Thresholding strategy: For each InterPro code carried by at least 3 proteins, the method sweeps 100 thresholds from 0 to the feature maximum and chooses the threshold maximizing F1.
Sampling strategy: For each SAE feature, the sample includes up to 60 proteins, combining the top 20 activators with proteins from four normalized-activation bins.

C Binarization of Continuous Annotation Methods

Continuous annotation methods convert geometry or sequence scores into binary activation labels for evaluation, using cross-validated or percentile-based thresholds.

Geometry method: GBM geometry scores use a threshold selected through 5-fold cross-validation to maximize F1 for residue-level classification metrics.
Sequence method: MEME binarizes residue activations at the 80th percentile over valid positions, labeling all nonzero activations positive when sparsity collapses that percentile to zero.

D Geometric annotation methods

The annotation pipeline combines protein-level, sequence-position, amino-acid, and multiscale geometric descriptors to predict where SAE features activate. Geometric descriptors quantify local backbone shape, contacts, and alignment across windows, while permutation-based evaluation tests feature associations.

Geometric descriptors: Geometric inputs describe Cα-backbone curvature, torsion, planarity, end-to-end shape, tangential alignment, and contact density.Curvature and torsion characterize local backbone geometry; planarity uses covariance eigenvalues, while contact density counts nearby or long-range outside residues.
Geometric descriptors: Multiscale statistics summarize means, maxima, and standard deviations of curvature, torsion, and planarity within local windows.Windows are split into thirds and compared at h/2 and 2h scales to reduce dependence on a single motif-relevant window size.
Geometric descriptors: Contact descriptors exclude residues inside the chosen window and measure outside residues within specified spatial and sequence-distance thresholds.The reported settings use X = 8, 12 and R = 12.
Covariates: Amino-acid covariates measure local fractions of residues grouped by flexibility, hydrophobicity, polarity, and aromaticity-related structural behavior.The groups include AGP, hydrophobic, apolar, and aromatic sets.
Baseline annotation methods: Protein-level InterPro methods test whether proteins activating on an SAE feature carry known database annotations, without locating activation within proteins.Their null hypothesis is evaluated by shuffling the vector of sampled proteins’ maximum activations.

G Geometric Feature Importance Breakdown

Most geometric annotations with PR-AUC > 0.3 show minimal dependence on amino-acid covariates, based on descriptor counts using GBM importance above 0.1.

Geometric Feature Importance Breakdown: The vast majority of geometric annotations show minimal dependence on amino-acid covariates.Figure 6 counts top descriptors with GBM importance above 0.1 among geometric annotations achieving PR-AUC > 0.3.

H Extension of runs to ESM-2 35M, layer 6

The analysis generalizes broadly from ESM-2 8M layer 4 to ESM-2 35M layer 6. Annotation coverage and PR-AUC distributions are broadly consistent between the models.

Extension to ESM-2 35M: Layer 6 of ESM-2 35M was selected as a middle-layer comparison with layer 4 of the eight-layer ESM-2 8M model.The 35M model has 12 layers, making layer 6 its middle layer.
Extension to ESM-2 35M: 92.99% of layer-6 features are annotated overall, while geometric annotation covers 92.92%.The same table reports 65.68% InterPro protein, 64.28% InterPro residue, 81.52% sequence-position, and 70.08% sequence-motif coverage.
Evaluation criterion: Feature annotation is defined by an FDR-controlled q value below 0.05 for at least one annotation method.The table defines overall annotation as satisfying this threshold for any method.

I Validation with Experimentally Determined Protein Structures

Validation on seven proteins with experimentally determined structures produced PR-AUC values closely matched to those obtained from AlphaFold structures. This suggests the geometric results are not driven by AlphaFold models.

Experimental-structure validation: PR-AUC values from AlphaFold and experimentally determined structures were closely matched across seven proteins.The comparison used layer 4 of ESM-2 8M; RMSD measured structural differences and nPos counted active residues for each protein-feature pair.