mlsanityData to 14 Oct 2026

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

arXiv 2606.21876 ↗First posted 20 Jun 2026cs.LG cs.AIcited by 0

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

  1. 2019What Does BERT Look At? An Analysis of BERT's Attentioncited by 1,003
  2. 2019A Multiscale Visualization of Attention in the Transformer Modelcited by 292
  3. 2019Evaluating Protein Transfer Learning with TAPEcited by 231
  4. 2020BERTology Meets Biology: Interpreting Attention in Protein Language Modelscited by 85
  5. 2024InterPLM: Discovering Interpretable Features in Protein Language Models via Sparse Autoencoderscited by 66
  6. 2019Learning protein sequence embeddings using information from structurecited by 51
  7. 2020What is a meaningful representation of protein sequences?cited by 22

Built on by

Not cited yet.