Protein contacts are already in the attention: a single-forward-pass alternative to the Categorical Jacobian
Not on the map. Why?
Abstract (arXiv)
The Categorical Jacobian of Zhang et al. (2024) reads protein contacts from a language model by perturbing every residue with every alternative amino acid, about 19L forward passes. We show the signal it reconstructs is already concentrated in a small subset of attention heads: averaging the top-K contact-relevant heads -- selected on as few as 10 labeled proteins, with no fitted per-pair or per-head weights -- recovers contacts in a single forward pass and matches or beats the Categorical Jacobian for every bidirectional model where it is defined (bar the smallest, 8M). Our primary test is leakage-clean: on a CAMEO split where neither selection nor evaluation touches data the models have plausibly memorized, the head readout beats the Categorical Jacobian on ESM-2-650M by +9pp (N = 29, p < 0.001), with the within-model margin reproducing across architectures. Ablations localize the gain to labeled head selection, not to averaging: at a matched label budget the unweighted mean ties a supervised L1 logistic regression on the same heads. Both methods fall 30-36pp from their in-distribution Zhang numbers to the leakage-clean split, which we read as an upper bound on how much prior numbers reflect pretraining overlap. We additionally introduce representation-CJ, a hidden-state generalization of the Jacobian to architectures without a masked-LM head (the output-head-independent analogue of logit-CJ), agreeing with the Categorical Jacobian where both are defined (per-protein Pearson r ≈ 0.95); show that the optimal K tracks how diffusely a model spreads its contact heads; and find both methods lose the signal on the two causal LMs we test, suggesting attention-encoded pair structure may depend on bidirectional pretraining.
Key facts. No citing papers in the map yet (data through 14 Oct 2026). Builds on 7 of its 30 references that are in the map. First posted on arXiv 20 Jun 2026.
Builds on
7 of 30 references.
References · 7
- 2019What Does BERT Look At? An Analysis of BERT's Attentioncited by 1,003
- 2019A Multiscale Visualization of Attention in the Transformer Modelcited by 292
- 2019Evaluating Protein Transfer Learning with TAPEcited by 231
- 2020BERTology Meets Biology: Interpreting Attention in Protein Language Modelscited by 85
- 2024InterPLM: Discovering Interpretable Features in Protein Language Models via Sparse Autoencoderscited by 66
- 2019Learning protein sequence embeddings using information from structurecited by 51
- 2020What is a meaningful representation of protein sequences?cited by 22
Built on by
Not cited yet.