Abstract:Protein language models (pLMs) have emerged as powerful predictors of protein structure and function. However, the computational circuits underlying their predictions remain poorly understood. Recent mechanistic interpretability methods decompose pLM representations into interpretable features, but they treat each layer independently and thus fail to capture cross-layer computation, limiting their ability to approximate the full model. We introduce ProtoMech, a framework for discovering computational circuits in pLMs using cross-layer transcoders that learn sparse latent representations jointly across layers to capture the model's full computational circuitry. Applied to the pLM ESM2, ProtoMech recovers 82-89% of the original performance on protein family classification and function prediction tasks. ProtoMech then identifies compressed circuits that use <1% of the latent space while retaining up to 79% of model accuracy, revealing correspondence with structural and functional motifs, including binding, signaling, and stability. Steering along these circuits enables high-fitness protein design, surpassing baseline methods in more than 70% of cases. These results establish ProtoMech as a principled framework for protein circuit tracing.
| Comments: | Accepted into ICML 2026. 32 pages, 17 figures |
| Subjects: | Machine Learning (cs.LG); Quantitative Methods (q-bio.QM) |
| Cite as: | arXiv:2602.12026 [cs.LG] |
| (or arXiv:2602.12026v2 [cs.LG] for this version) | |
| https://doi.org/10.48550/arXiv.2602.12026 arXiv-issued DOI via DataCite |
Submission history
From: Darin Tsui [view email]
[v1]
Thu, 12 Feb 2026 14:57:57 UTC (22,074 KB)
[v2]
Wed, 13 May 2026 01:22:30 UTC (19,507 KB)