Mechanistic evidence that motif-gated domain recognition drives contact prediction in protein language models
Nainani, J.; Reimer, B. M.; Watts, C.; Jensen, D.; Green, A. G.
Show abstract
Protein language models (pLMs) achieve state-of-the-art performance on protein structure and function prediction tasks, yet their internal computations re-main opaque. Sparse autoencoders (SAEs) have been used to recover sparse features, called latents, from pLM layer representations, whose activations cor-relate with known biological concepts. However, prior work has not established which model concepts are causally necessary for pLM performance on down-stream tasks. Here, we adapt causal activation patching to the pLM setting and perform it in SAE latent space to extract the minimal circuit responsible for accuracy in a contact prediction task for two case study proteins. We observe that preserving only a tiny fraction of latent-token pairs (0.022% and 0.015%) is sufficient to retain contact prediction accuracy in a residue unmasking experiment. Our circuit indicates a two-step computation in which early-layer motif detectors respond to short local sequence patterns, gating mid-to-late domain detectors which are selective for protein domains and families. Path-level ab-lations confirm the causal dependence of domain latents on upstream motif latents. To evaluate these components quantitatively, we introduce two diagnostics: a Motif Conservation Test and a Domain Selectivity Framework that supports hypothesis-driven tests. All candidate motif-detector latents pass the conservation test, and 18/23 candidate domain-detector latents achieve AUROC [≥]0.95. To our knowledge, this is the first circuits-style causal analysis for pLMs, pin-pointing the motifs, domains, and motif-domain interactions that drive contact prediction in two specific case studies. The framework introduced herein will enable future mechanistic dissection of protein language models. Code available at https://github.com/NainaniJatinZ/plm_circuits
Matching journals
The top 5 journals account for 50% of the predicted probability mass.
Similar papers in this journal
Similar papers in this journal
- Discovering molecular features of intrinsically disordered regions by using evolution for contrastive learning 94%
- Deep Mendelian Randomization: Investigating the causal knowledge of genomic deep learning models 94%
- Combining phylogeny and coevolution improves the inference of interaction partners among paralogous proteins 94%
"Similar papers" are the closest papers from that journal in the model's embedding space. They show what the match is built on, but the ranking comes mostly from a classifier over the whole training set, not from these examples alone.