arXiv:2606.21876cs.LGcs.AI2026-06

用少量蛋白数据选注意力头,单次前向传播即可读取蛋白质接触信息。

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

论文配图:Protein contacts are already in the attention: a single-forward-pass alternative to the Categorical Jacobian
图 1 · 摘自论文原文
  • 仅通过10个标注蛋白选择关键注意力头,无需额外参数即可复现接触信号。
  • 单次前向传播性能超越传统方法,在清洁测试集上提升9个百分点(p<0.001)。
  • 适用于不依赖输出头的模型,且对不同架构具有通用性,适合结构预测研究者。

Zhang等人(2024)的类别雅可比方法通过扰动每个残基为所有替代氨基酸,需约19L次前向传播来读取蛋白质接触信息。我们发现该信号已集中在少数注意力头中:仅平均前K个与接触相关的头部——基于少至10个标注蛋白选择,且无需每对或每头的拟合权重——即可在单次前向传播中恢复接触信息,并在所有定义的双向模型中达到或超过类别雅可比表现(除最小的8M模型外)。主要测试采用无泄漏分割:在未被模型记忆的数据上,该方法在ESM-2-650M上比类别雅可比高出9个百分点(N=29,p<0.001),且模型间差异一致。消融实验表明增益来自有标签头的选择,而非平均:在相同标注预算下,未加权平均与监督式L1逻辑回归性能相当。两种方法在清洁分割上均较原内分布结果下降30-36个百分点,这提示先前数值可能受预训练重叠影响。我们还引入表示型雅可比(representation-CJ),将其扩展至无掩码语言模型头的架构(即对隐藏状态的雅可比类比),在两者都适用时与类别雅可比高度一致(每蛋白皮尔逊相关系数r≈0.95);并发现最优K值反映模型接触头分布的扩散程度;最后发现两类方法在两个因果语言模型上均失效,暗示注意力编码的配对结构可能依赖双向预训练。

原文摘要 · Abstract (English)

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.

蛋白质结构注意力机制模型解释深度学习

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。