View PDF HTML (experimental)

Abstract:The Categorical Jacobian (CJ) 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, recovers contacts in one forward pass and beats CJ on leakage-clean data for every bidirectional model where CJ is defined, and matches or beats it in-distribution (the exceptions being the smallest 8M model and a statistical tie on ESM Cambrian). Ablations localize the gain to labeled head selection, not averaging: at a matched label budget the unweighted mean ties a supervised L1 logistic regression on the same heads, so the parameter-free mean is selection's minimal form, not the source of the advantage. 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 CJ on ESM-2-650M by +9 pp (N=29, p<0.001), with the within-model margin reproducing across architectures on a wider pretraining-aware set. Both methods fall 30-36 percentage points from their in-distribution Zhang numbers to the leakage-clean numbers, consistent with substantial pretraining overlap inflating prior numbers (a CAMEO-vs-Zhang difficulty shift contributes too, so we read it as an upper bound on the leakage component). We additionally introduce representation-CJ, a hidden-state generalization of the Jacobian for architectures without a masked-LM head; show that the optimal K tracks how diffusely a model spreads its contact heads; and find that both methods lose the contact signal on both causal LMs we test (ProGen2), suggesting attention-encoded pair structure may depend on bidirectional pretraining.

Submission history

From: Rome Thorstenson [view email]
[v1] Sat, 20 Jun 2026 04:35:51 UTC (96 KB)