# Protein contacts are already in the attention: a single-forward-pass alternative to the Categorical Jacobian

arXiv: https://arxiv.org/abs/2606.21876 · HTML: https://mlsanity.com/paper/2606.21876/

## 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 (most cited first, 7 of 7)

- [What Does BERT Look At? An Analysis of BERT's Attention](https://mlsanity.com/paper/1906.04341/) (2019)
- [A Multiscale Visualization of Attention in the Transformer Model](https://mlsanity.com/paper/1906.05714/) (2019)
- [Evaluating Protein Transfer Learning with TAPE](https://mlsanity.com/paper/1906.08230/) (2019)
- [BERTology Meets Biology: Interpreting Attention in Protein Language Models](https://mlsanity.com/paper/2006.15222/) (2020)
- [InterPLM: Discovering Interpretable Features in Protein Language Models via Sparse Autoencoders](https://mlsanity.com/paper/2412.12101/) (2024)
- [Learning protein sequence embeddings using information from structure](https://mlsanity.com/paper/1902.08661/) (2019)
- [What is a meaningful representation of protein sequences?](https://mlsanity.com/paper/2012.02679/) (2020)

## 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.