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 L_1 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 \approx 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.
Protein contacts are already in the attention: a single-forward-pass alternative to the Categorical Jacobian
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.
- Preview

- Year
- 2026
- Hosting
- Full text hostedCC-BY-4.0
Cite
Notes
Only stored in your browser.
Attribution
- Abstract & full text
- arxiv.org/abs/2606.21876CC-BY-4.0
- TL;DR
- Semantic Scholar