🤖 machine learning

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

本文表明，此前通过计算成本高昂的、需要约 19 层前向传播的分类雅可比（Categorical Jacobian）提取出的蛋白质接触信号，已经集中在可以用极少量标记数据识别出的少量注意力头中，且在无泄漏基准测试上，仅需单次前向传播即可实现更优或相当的性能。

原作者： Rome Thorstenson

发布于 2026-06-23

📖 1 分钟阅读☕ 轻松阅读

CC BY 4.0

原作者： Rome Thorstenson

原始论文采用 CC BY 4.0 许可（http://creativecommons.org/licenses/by/4.0/）。 ✨ 这是对下方论文的AI生成解释。它不是由作者撰写或认可的。如需技术准确性，请参阅原始论文。阅读完整免责声明

核心思想：寻找蛋白质“大脑”中的捷径

想象一下，蛋白质语言模型（PLM）就像一名超级聪明的学生，读过数百万条蛋白质序列。因此，这位学生已经学会了蛋白质如何折叠成三维形状，尽管他从未被明确教授过几何学知识。

科学家们一直试图“读懂”这位学生的思想，看看蛋白质的哪些部分会粘在一起（接触点），从而形成那个三维形状。

旧方法（“分类雅可比矩阵”/ Categorical Jacobian）：
以前最好的方法就像是一个非常认真但极其劳累的侦探。为了弄清楚蛋白质的两个部分是否是“好朋友”，侦探会：

取出蛋白质。
将序列的第一个字母改为其他每一种可能的字母（共19种）。
观察学生的答案发生了怎样的变化。
对蛋白质中的每一个字母重复上述步骤。

如果一个蛋白质有300个字母，这个侦探为了得到一个答案，必须让学生进行 5,700 次测试（300个字母 × 19次改变）。这种方法很准确，但速度极慢，计算成本极高。

新方法（“注意力融合”/ Attention Fusion）：
本文认为，这位学生在第一次阅读蛋白质时，其实已经在笔记中写下了答案。这些笔记被称为“注意力图”（attention maps）。

作者发现，如果你查看学生的笔记，关于蛋白质哪些部分会“粘在一起”的信息已经存在于其中，并且集中在仅仅几个特定的“头”（heads）里（可以把这些想象成学生大脑中特定的笔记本或页面）。

与其运行5,700次测试，新方法只需：

让蛋白质通过学生进行一次运行。
查看学生写下“粘连”笔记的特定页面（注意力头）。
将这些页面进行平均，从而得到答案。

结果： 新方法与旧的侦探方法一样准确（甚至往往更好），但它只需要 一次运行，而不是数千次。这就像是在教科书背面找到了“小抄”，而不是从头开始重新推导整个公式。

关键发现详解

1. “泄露”问题（作弊测试）

作者注意到，之前的研究可能存在“作弊”行为。用于评估这些模型的测试蛋白质，有时与学生在训练期间已经背过的蛋白质非常相似。这就像是给一个学生参加期末考试，而考试题目竟然出现在之前的家庭作业里。

为了解决这个问题，作者创建了一个“无泄露”测试。他们只使用了在学生完成学习之后才发现的蛋白质。

发生了什么？ 当他们去除了这种“作弊”优势后，旧的侦探方法（分类雅可比矩阵）表现得大打折扣。而新的“小抄”方法（注意力融合）依然保持强劲。
结论： 旧方法在一定程度上依赖于学生对特定测试题目的记忆。而新方法则真正理解了结构。

2. “弥散型” vs. “集中型”笔记本

作者发现不同的模型，其记录笔记的方式也不同：

集中型模型（Concentrated Models）： 一些模型将答案写在一个特定的笔记本里。如果你看那一个笔记本，就能得到答案。
弥散型模型（Diffuse Models）： 其他模型则将答案分散在许多个笔记本中。没有任何一个笔记本拥有完整的图景，但如果你把它们全部平均起来，答案就会显现。

论文表明，你可以通过观察几个例子来判断你拥有哪种类型的模型，然后你就知道该看一个笔记本，还是平均十个笔记本。

3. “因果型”死胡同（The "Causal" Dead End）

作者测试了另一种类型的模型（“因果型”模型），这种模型像读句子一样从左到右阅读蛋白质，且不会向后看。

结果： 无论是旧的侦探方法还是新的小抄方法，在这些模型上都完全失败了。
结论： 这表明，要理解蛋白质如何在三维空间中粘合在一起，模型需要能够同时观察整个蛋白质（双向性），而不仅仅是逐字逐句地阅读。

4. 速度与成本

两者在速度上的差异是巨大的。

旧方法： 如果你想分析1,000个蛋白质，旧方法可能会消耗数百美元的计算时间，因为它必须对每个蛋白质进行数千次模拟。
新方法： 它的成本仅为一小部分，因为它每个蛋白质只需运行一次。这就像是手动数清沙滩上的每一粒沙子，与拍摄一张卫星照片并使用简单算法估算总数之间的区别。

总结

本文声称，蛋白质模型已经知道蛋白质是如何折叠的；它们只是将这些知识隐藏在内部的“注意力”机制中。通过学习直接读取这些特定的笔记，我们可以比以前更快、更准确地预测蛋白质结构，而无需进行数千次昂贵的模拟。这种方法在那些能够同时观察整个蛋白质的模型上效果最好，并且即使在测试模型从未见过的全新数据时，依然表现稳健。

技术摘要：蛋白质接触信息已存在于注意力机制中

问题陈述

在海量无标签序列语料库上训练的蛋白质语言模型（PLMs）学习到了丰富的蛋白质结构表示。一个核心问题是，这种结构知识是否可以在无需结构标签监督训练的情况下被提取出来。Zhang 等人 [2024] 提出了**标准分类雅可比矩阵（Standard Categorical Jacobian, CJ）**作为解决方案：通过对每个残基进行每种替代氨基酸的扰动，并测量模型在掩码语言模型（MLM）逻辑值（logits）上的响应，可以重建接触图。虽然 CJ 实现了较高的精确度（在某些基准测试中 P@L ≈ 0.80），但其计算成本高昂，对于长度为 $L$ 的蛋白质，大约需要 $19L$ 次前向传播。此外，CJ 在结构上受限于具有每位置分类头的架构（双向掩码语言模型），这限制了其在编码器-解码器或因果模型中的适用性。

近期的评估也引发了对数据泄漏的担忧，即预训练语料库可能与评估集存在重叠，从而夸大了性能指标。因此，需要一种计算高效、适用于多种架构且对预训练泄漏具有鲁棒性的方法。

方法论

作者提出，由分类雅可比矩阵重建的接触信号已经显式地存在于模型的**注意力图（attention maps）**中，并集中在少数可识别的注意力头子集中。

1. Top-K 注意力融合

该方法不再通过扰动输入，而是直接从 PLM 的注意力机制中提取接触信号：

提取： 对于每个注意力头，对注意力图进行对称化、去对角线处理，并使用平均乘积校正（Average Product Correction, APC）进行修正。
选择： 通过在小型、不相交的有标签选择集（仅需 10 个蛋白质）上衡量其预测接触的能力（以 top-L/2 长程精确度衡量），对注意力头进行排名。
融合： 将选出的前 $K$ 个注意力头通过简单的无权重算术平均值进行组合。在推理阶段，这仅需要针对每个蛋白质进行一次前向传播。
标签效率： 前 $K$ 个头的选择是唯一的监督步骤，仅需极少的有标签蛋白质预算（用于选择的 10 个，以及用于与监督回归进行匹配预算对照的 50 个）。

2. 表示-CJ (Representation-CJ)

为了在缺乏 MLM 头（如编码器-解码器或因果模型）的架构上与 CJ 进行比较，作者引入了 Representation-CJ。它通过将逻辑响应替换为隐藏状态响应来推广雅可比构造：
$J_{repr}[i, a', j] = \| h_{i \to a'}[\ell, j] - h_{unpert}[\ell, j] \|_2$
其中 $h$ 代表当残基 $i$ 被扰动为氨基酸 $a'$ 时，位置 $j$ 处的层 $\ell$ 的隐藏状态。这保持了与 logit-CJ 相同的计算成本（ $19L + 1$ 次前向传播），但适用于任何暴露隐藏状态的架构。

3. 评估协议

研究采用了**防泄漏（leakage-clean）**评估策略，以确保结果不是预训练重叠的产物：

CAMEO-PTA25： 主要评估集由在模型预训练截止日期（2022-09）之后提交的蛋白质链组成，并针对当前 UniRef50 以 50% 的恒等度进行了过滤。
不相交选择： 头选择和评估在不相交的数据子集上执行，以防止信息泄漏。
基准测试： 该方法与标准分类雅可比矩阵（logit-CJ）、最佳单个注意力头以及监督 L1-逻辑回归基准进行了比较。

核心贡献

效率与性能： Top-K 注意力融合方法通过单次前向传播即可恢复接触信号，在所有测试的双向模型（ESM-2 系列、ESM-1b、ProtT5-XL）上，其表现均优于分类雅可比矩阵。
防泄漏验证： 在防泄漏的 CAMEO 分割集上，融合方法的表现显著优于 CJ（例如在 ESM-2-650M 上提升了 9 个百分点）。至关重要的是，在这些干净的数据上，分类雅可比矩阵不再优于最佳单个注意力头，这表明其之前的分布内（in-distribution）优越性部分源于记忆效应。
架构泛化性： 作者证明了接触信号可以通过注意力融合存在于编码器-解码器模型（ProtT5-XL）中，并引入了 Repr-CJ 以将基于雅可比的分析扩展到这些架构。
机制洞察（弥散性）： 融合最佳注意力头的数量 ( $K$ $K$ ) 取决于模型内接触信号的“弥散性（diffuseness）”：
- 集中簇（Concentrated Clusters）： 一些模型（如 ESM-2-8M、AMPLIFY-350M）将信号集中在一个主导头上；增加平均更多的头会稀释信号。
- 弥散簇（Diffuse Clusters）： 更大的模型（如 ESM-2-650M、ProtT5-XL）将信号分散在许多头中；需要通过平均化来恢复共识。
因果模型局限性： 融合法和 CJ 方法在因果语言模型（ProGen2-large 和 -xlarge）上均失效，这表明双向预训练对于注意力机制编码对称对结构是必要的。

结果

性能差距： 在防泄漏数据上，融合法在 ESM-2-650M 上比 CJ 高出约 9 个百分点。当从分布内（Zhang eval-200）移动到防泄漏（CAMEO-PTA25）基准测试时，两种方法在绝对精确度上都下降了 30–36 个百分点，这表明先前的无监督接触预测结果由于预训练重叠而被显著夸大了。
监督 vs. 无监督： 在匹配的标签预算（50 个蛋白质）下，前 $K$ 个头的无权重平均值在统计上与基于相同头的监督 L1-逻辑回归拟合表现无异。这意味着注意力的选择是性能的主要驱动力，而非读出方式的参数形式。
计算经济性： 融合法每个蛋白质仅需一次前向传播，而 CJ 需要 $\sim 19L$ 次。计算成本的盈亏平衡点在 20 个测试蛋白质以下，使得融合法在大规模应用中更具效率。
Repr-CJ 验证： 在两者皆定义的架构上，Repr-CJ 与 logit-CJ 高度相关（Pearson $r > 0.90$ ），验证了其作为非 MLM 模型忠实泛化形式的有效性。

意义与主张

本文声称，分类雅可比矩阵并非提取了新的接触信息，而是重建了模型在其注意力模式中已经表示的结构对信息。主要发现是，这种信号可以通过对一小部分注意力头进行简单的、无参数的平均来获取，从而提供了一种计算上更优的替代方案。

作者强调，“防泄漏”评估揭示了分布内性能与真实泛化性能之间的巨大差距，这表明未来的无监督接触预测基准必须默认采用考虑预训练情况的划分方式。这项工作还强调，结构信息的分配是在注意力头之间进行的学习属性，而非架构上的必然性，即使在同一模型家族内（例如 ESM-2 的不同规模）也会有所不同。

最后，两种方法在因果语言模型上的失效被视为证据，表明双向预训练是注意力机制学习对称对结构的批判性因素，尽管作者指出鉴于测试的因果模型范围有限，这是一个关联性发现。