🤖 machine learning

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

This paper demonstrates that protein contact signals, previously extracted via the computationally expensive Categorical Jacobian requiring ~19L forward passes, are already concentrated in a small subset of attention heads that can be identified with minimal labeled data to achieve superior or comparable performance in a single forward pass on leakage-clean benchmarks.

Original authors: Rome Thorstenson

Published 2026-06-23

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Rome Thorstenson

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

The Big Idea: Finding a Shortcut in the Protein "Brain"

Imagine a protein language model (PLM) as a super-smart student who has read millions of protein sequences. Because of this, the student has learned how proteins fold into 3D shapes, even though they were never explicitly taught geometry.

Scientists have been trying to "read the mind" of this student to see which parts of the protein stick together (contacts) to form that 3D shape.

The Old Way (The "Categorical Jacobian"):
Previously, the best method was like a very thorough but exhausting detective. To figure out if two parts of a protein are friends, the detective would:

Take the protein.
Change the first letter of the sequence to every other possible letter (19 times).
Watch how the student's answer changes.
Repeat this for every single letter in the protein.

If a protein has 300 letters, this detective has to run the student through the test 5,700 times (300 letters × 19 changes) just to get one answer. It's accurate, but it's incredibly slow and computationally expensive.

The New Way (The "Attention Fusion"):
This paper argues that the student actually wrote the answer down in their notes while they were reading the protein the first time. These notes are called "attention maps."

The authors found that if you look at the student's notes, the information about which parts of the protein stick together is already there, concentrated in just a few specific "heads" (think of these as specific notebooks or pages in the student's brain).

Instead of running the test 5,700 times, the new method simply:

Runs the protein through the student once.
Looks at the specific pages (attention heads) where the student wrote down the "sticking together" notes.
Averages those pages together to get the answer.

The Result: The new method is just as accurate (and often better) than the old detective method, but it takes one single pass instead of thousands. It's like finding a cheat sheet in the back of the textbook instead of re-deriving the whole formula from scratch.

Key Findings Explained

1. The "Leakage" Problem (The Cheating Test)

The authors noticed that previous studies might have been "cheating." The test proteins used to evaluate these models were sometimes similar to the proteins the student had already memorized during training. It's like giving a student a final exam that uses questions they already saw in the homework.

To fix this, the authors created a "leakage-clean" test. They used only proteins that were discovered after the student finished studying.

What happened? When they removed the "cheating" advantage, the old detective method (Categorical Jacobian) got much worse. The new "cheat sheet" method (Attention Fusion) stayed strong.
The Takeaway: The old method was partly relying on the student's memory of the specific test questions. The new method actually understands the structure better.

2. The "Diffuse" vs. "Concentrated" Notebooks

The authors discovered that different models write their notes differently:

Concentrated Models: Some models write the answer in just one specific notebook. If you look at that one notebook, you get the answer.
Diffuse Models: Other models spread the answer across many notebooks. No single notebook has the full picture, but if you average them all together, the answer appears.

The paper shows that you can figure out which type of model you have by looking at a few examples, and then you know whether to look at one notebook or average ten.

3. The "Causal" Dead End

The authors tested a different type of model (a "Causal" model) that reads proteins like a sentence, from left to right, without looking back.

The Result: Both the old detective method and the new cheat sheet method failed completely on these models.
The Takeaway: This suggests that to understand how proteins stick together in 3D, the model needs to be able to look at the whole protein at once (bidirectional), not just read it word-by-word.

4. Speed and Cost

The difference in speed is massive.

Old Method: If you want to analyze 1,000 proteins, the old method might cost you hundreds of dollars in computer time because it has to run the simulation thousands of times per protein.
New Method: It costs a fraction of that because it only runs once per protein. It's like the difference between manually counting every grain of sand on a beach versus taking a satellite photo and using a simple algorithm to estimate the total.

Summary

The paper claims that protein models already know how proteins fold; they just hide that knowledge in their internal "attention" mechanisms. By learning to read those specific notes directly, we can predict protein structures much faster and more accurately than before, without needing to run thousands of expensive simulations. This works best on models that can look at the whole protein at once, and it holds up even when we test the models on brand-new data they haven't seen before.

Technical Summary: Protein Contacts are Already in the Attention

Problem Statement

Protein Language Models (PLMs) trained on vast unlabeled sequence corpora learn rich representations of protein structure. A central question is whether this structural knowledge can be extracted without supervised training on structure labels. Zhang et al. [2024] proposed the Standard Categorical Jacobian (CJ) as a solution: by perturbing every residue with every alternative amino acid and measuring the response in the model's masked-language-model (MLM) logits, one can reconstruct a contact map. While CJ achieves high precision (P@L ≈ 0.80 on certain benchmarks), it is computationally expensive, requiring approximately $19L$ forward passes for a protein of length $L$ . Furthermore, CJ is structurally tied to architectures with a per-position categorical head (bidirectional masked-LMs), limiting its applicability to encoder-decoder or causal models.

Recent evaluations have also raised concerns about data leakage, where pretraining corpora may overlap with evaluation sets, inflating performance metrics. There is a need for a method that is computationally efficient, applicable across diverse architectures, and robust to pretraining leakage.

Methodology

The authors propose that the contact signal reconstructed by the Categorical Jacobian is already explicitly present in the model's attention maps, concentrated in a small, identifiable subset of attention heads.

1. Top-K Attention Fusion

Instead of perturbing inputs, the proposed method extracts contact signals directly from the attention mechanisms of the PLM:

Extraction: For each attention head, the attention map is symmetrized, zero-diagonalized, and corrected using Average Product Correction (APC).
Selection: Heads are ranked by their ability to predict contacts (measured by top-L/2 long-range precision) on a small, disjoint labeled selection set (as few as 10 proteins).
Fusion: The top- $K$ selected heads are combined using a simple, unweighted arithmetic mean. This requires only one forward pass per protein at inference time.
Label Efficiency: The selection of the top- $K$ heads is the only supervised step, requiring a minimal budget of labeled proteins (10 for selection, 50 for a matched-budget control against supervised regression).

2. Representation-CJ (Repr-CJ)

To compare against CJ on architectures lacking an MLM head (e.g., encoder-decoder or causal models), the authors introduce Representation-CJ. This generalizes the Jacobian construction by replacing the logit response with the hidden-state response:
$J_{repr}[i, a', j] = \| h_{i \to a'}[\ell, j] - h_{unpert}[\ell, j] \|_2$
where $h$ represents the hidden state at position $j$ and layer $\ell$ when residue $i$ is perturbed to amino acid $a'$ . This maintains the same computational cost ( $19L + 1$ forward passes) as logit-CJ but applies to any architecture exposing hidden states.

3. Evaluation Protocol

The study employs a leakage-clean evaluation strategy to ensure results are not artifacts of pretraining overlap:

CAMEO-PTA25: A primary evaluation set consisting of protein chains deposited after the pretraining cutoff of the models (post-2022-09), filtered against current UniRef50 at 50% identity.
Disjoint Selection: Head selection and evaluation are performed on disjoint subsets of the data to prevent information leakage.
Baselines: The method is compared against the Standard Categorical Jacobian (logit-CJ), the best single attention head, and supervised L1-logistic regression baselines.

Key Contributions

Efficiency and Performance: The Top-K attention fusion method recovers contact signals in a single forward pass, beating the Categorical Jacobian on leakage-clean data for every bidirectional model tested (ESM-2 family, ESM-1b, ProtT5-XL).
Leakage-Clean Validation: On a leakage-clean CAMEO split, the fusion method significantly outperforms CJ (e.g., +9 percentage points on ESM-2-650M). Crucially, on this clean data, the Categorical Jacobian no longer outperforms the best single attention head, suggesting its previous in-distribution superiority was partly due to memorization.
Architectural Generalization: The authors demonstrate that the contact signal exists in encoder-decoder models (ProtT5-XL) via attention fusion and introduce Repr-CJ to extend Jacobian-based analysis to these architectures.
Mechanistic Insight (Diffuseness): The optimal number of heads ( $K$ $K$ ) to fuse depends on the "diffuseness" of the contact signal within the model.
- Concentrated Clusters: Some models (e.g., ESM-2-8M, AMPLIFY-350M) concentrate the signal in a single dominant head; averaging more heads dilutes the signal.
- Diffuse Clusters: Larger models (e.g., ESM-2-650M, ProtT5-XL) spread the signal across many heads; averaging is required to recover the consensus.
Causal Model Limitation: Both fusion and CJ methods collapse on causal language models (ProGen2-large and -xlarge), suggesting that bidirectional pretraining is necessary for attention mechanisms to encode symmetric pair structures.

Results

Performance Gap: On leakage-clean data, fusion beats CJ by approximately 9 percentage points on ESM-2-650M. Both methods drop 30–36 percentage points in absolute precision when moving from in-distribution (Zhang eval-200) to leakage-clean (CAMEO-PTA25) benchmarks, indicating that prior unsupervised contact prediction results were significantly inflated by pretraining overlap.
Supervised vs. Unsupervised: At a matched label budget (50 proteins), an unweighted mean of the top- $K$ heads performs statistically indistinguishably from a supervised L1-logistic regression fit on the same heads. This implies that the selection of heads is the primary driver of performance, not the parametric form of the readout.
Compute Economy: Fusion requires one forward pass per protein, whereas CJ requires $\sim 19L$ . The break-even point for computational cost is below 20 test proteins, making fusion vastly more efficient for large-scale applications.
Repr-CJ Validation: Repr-CJ correlates tightly with logit-CJ (Pearson $r > 0.90$ ) on architectures where both are defined, validating it as a faithful generalization for non-MLM models.

Significance and Claims

The paper claims that the Categorical Jacobian does not extract new contact information but rather reconstructs structural pair information that the model already represents in its attention patterns. The primary finding is that this signal is accessible via a simple, parameter-free averaging of a small subset of attention heads, offering a computationally superior alternative to perturbation-based methods.

The authors emphasize that the "leakage-clean" evaluation reveals a substantial gap between in-distribution and true generalization performance, suggesting that future benchmarks for unsupervised contact prediction must adopt pretraining-aware splits as a default. The work also highlights that the allocation of structural information across attention heads is a learned property of the pretraining process rather than an architectural inevitability, varying even within the same model family (e.g., ESM-2 scales).

Finally, the collapse of both methods on causal LMs is presented as evidence that bidirectional pretraining is a critical factor for learning symmetric pair structures in attention, though the authors note this is an associational finding given the limited scope of causal models tested.