通过可微图划分对蛋白质语言模型表示的结构解释

# 通过可微图划分对蛋白质语言模型表示的结构化解释

来源: https://arxiv.org/html/2605.10985

Edward Tan Beng Wai$^{1,\dagger}$, Soumick Sarker$^{1,\dagger}$, Pasan Gunawardane$^{1,\dagger}$, Jagath C. Rajapakse

新加坡南洋理工大学
{ siddhant010, soumick001, ed0001ai, c250135 }@e.ntu.edu.sg, [email protected]

###### 摘要

像 ESM-2 这样的蛋白质语言模型（Protein Language Models, PLMs）学习了丰富的残基表示，在蛋白质功能预测任务中取得了强大的性能，但其特征仍然难以解释，因为结构性和进化信号被编码在密集的空间中。我们提出了一种即插即用的框架，将 ESM-2 表示投影到蛋白质接触图上，并应用 SoftBlobGIN——一种轻量级的图同构网络（Graph Isomorphism Network），具有可微的 Gumbel-softmax 子结构池化功能——以执行结构感知的消息传递，并为下游预测任务学习粗粒度的功能子结构。在酶分类任务中，SoftBlobGIN 达到了 92.8% 的准确率和 0.898 的宏 F1 分数。与仅对蛋白质语言模型进行事后分析不同，我们的方法产生了可直接审计的结构化解释：GNNExplainer 恢复了具有生物学意义的活性位点残基、空间局部化的功能簇以及催化接触模式。在结合位点检测任务上，SoftBlobGIN 将残基 AUROC 从仅使用 ESM-2 线性探测器的 0.885 提升至 0.983，表明这些结构化解释无法仅从语言模型特征中恢复。学习到的“blob”分区通过自动将残基分组为功能子结构提供了额外的可解释性层次，含有注释活性位点残基的 blobs 比其它 blobs 具有高 1.85 倍的重要性（$\rho=0.339, p=0.009$），且无需任何活性位点监督。我们的框架不需要重新训练语言模型，仅增加约 110 万参数，并在 ProteinShake 任务上泛化良好，在基因本体预测中达到 $F_{max}$ 0.733，在结合位点检测中达到 AUROC 0.969。我们将此定位为蛋白质语言模型的可解释结构化伴侣，使其预测更加透明和可审计。

## 1 引言

尽管序列数据库激增，但仍有大量已测序的蛋白质未被功能注释（Kustatscher et al., 2022）。大规模结构数据（如 AlphaFold 2 (Jumper et al., 2021)）的兴起为超过 2 亿种蛋白质提供了预测结构。最近，蛋白质语言模型（PLMs）如 ESM-2 (Lin et al., 2023) 改变了功能预测的局面。特别是 ESM-2 已成为事实上的特征提取器：对其均值池化嵌入使用简单的多层感知机（MLPs）在许多任务上已经匹配或超越了基于结构的图神经网络（GNNs）。然而，这种成功是以**可解释性**为代价的。ESM-2 嵌入是 1280 维的密集向量，没有明显的映射到特定残基、接触或生化基序。然而，许多下游任务需要可解释性，特别是在安全和合规至关重要的临床环境中部署时。此外，可解释的模型允许研究人员验证计算预测是否对应于有意义的生物机制，而不是虚假的数据相关性，从而可能发现新的见解。

现有的语言模型探针（基于注意力头的线性分类器、注意力回滚等）恢复了广泛的序列模式，但很少揭示空间局部化、生化特定的基序。另一方面，结构方法（如基于蛋白质接触图的传统 GNN）通常使用固定半径的邻域，限制了经验定义常数的灵活性。这意味着与底物形成 3 Å 氢键的催化残基和距离远端环 7.9 Å 的表面残基会收到相同的图拓扑结构。最近的工作（Wang and Oliver, 2025）通过使用变量大小的分区来解决这个问题，基于几何向量感知器（GVP）编码器和向量量化（VQ）码本，但这以增加计算成本和可解释性差距为代价。重要的是，这些方法没有利用近期 PLM 表示的表达力。这引出了一个问题：何时**结构**推理能增加 PLM 已经捕获的信息之外的信息？

我们的经验答案是：边界在于生物学。对于图级功能任务如酶分类（EC），ESM-2 均值池化几乎足够了，图结构增加的见解很少。对于残基级结构任务如结合位点检测，在接触图上进行的消息传递增加了 ESM-2 单独无法恢复的大量信息。有趣的后果是可解释性：在结构重要的领域，我们希望模型的**结构推理**是可审计的。因此，我们提出了一种计算轻量级、结构可解释的 GNN，它联合利用了语义丰富的 ESM-2 特征，同时保持可解释性。为此，我们列出我们的贡献：

1.  **经验表征**：我们映射了何时结构推理有助于冻结的 PLM 特征。对于图级 EC 任务，ESM-2 均值池化几乎足够了（有无接触图的准确率分别为 0.910 vs 0.912）。对于残基级结合位点检测，在接触图上的消息传递弥补了 ESM-2 单独无法弥补的 9.8 分 AUROC 差距。
2.  **可解释架构**：我们引入了 SoftBlobGIN，用单个 Gumbel-softmax 分配头（Jang et al., 2016）替换 BioBlobs 的 GVP 编码器和 VQ 码本，产生 $K$ 个可微的、软化的蛋白质子结构，具有约 110 万个参数且无需重新训练语言模型。
3.  **生物学验证**：我们定量验证了结果解释与已建立的酶生化的一致性。GNNExplainer (Ying et al., 2019) 恢复了催化残基富集、活性位点埋藏、空间共定位以及与催化三联体一致的三级接触几何结构。学习到的 blobs 自发地将功能位点与结构支架分离，含有活性位点的 blobs 承载了高 1.85 倍的重要性（$\rho=0.339, p=0.009$）。

## 2 相关工作

#### 用于蛋白质结构的图神经网络
蛋白质结构自然地表示为残基接触图，其中节点对应于氨基酸，边编码空间邻近性或几何关系。传统的 GNN 如图卷积网络（GCN）(Kipf and Welling, 2016) 和图注意力网络（GAT）(Veličković et al., 2017) 因此被广泛用于基于结构的蛋白质学习。然而，这些架构的表达力有限。Xue et al. (2018a) 表明，GCN 和 GAT 严格来说比 Weisfeiler-Leman (WL) 图同构测试表达力弱，这促使了图同构网络（GIN）的引入，后者达到了 WL 上界。后续工作通过 GINEConv (Hu et al., 2019) 扩展了 GIN 以整合边信息，使得明确建模成对残基几何结构（如 C$\alpha$–C$\alpha$ 距离）成为可能。这些发展使 GIN 风格架构成为蛋白质图的自然基础，其中局部几何相互作用通常是功能信息丰富的。

#### 蛋白质语言模型（PLMs）
与此同时，PLM 通过大规模学习序列衍生的表示极大地推动了功能预测。像 ESM-2 (Lin et al., 2023) 这样的模型在约 6500 万种蛋白质序列上使用掩码语言建模进行训练，并产生丰富的每残基嵌入，编码进化保守性、结构规律性和功能上下文。在许多下游场景中，冻结的 ESM-2 特征结合轻量级分类器已经实现了具有竞争力或最先进的性能。这确立了 PLM 作为强通用特征提取器的地位，但也引入了可解释性挑战。因为这些表示是高维密集向量，负责预测的结构或生化信号并不直接可见。我们的工作基于这一观察，将 ESM-2 视为冻结的语义编码器，同时引入显式的基于图的结构推理。

#### 层次池化与结构抽象
除了残基级消息传递外，层次池化方法旨在通过将残基分组为更高层次的子结构来学习粗粒度的结构抽象。DiffPool (Ying et al., 2018) 引入了用于层次图粗化的可微软聚类分配，实现了图层次结构的端到端学习。最近，BioBlobs (Wang and Oliver, 2025) 通过受生物启发的 blob 分区将此范式适应于蛋白质，在 ProteinShake 基准测试中展示了强大的性能。然而，这种方法依赖于更重的架构，涉及几何向量感知器和向量量化。相比之下，我们的 SoftBlobGIN 用轻量级的 Gumbel-softmax 池化 (Jang et al., 2016) 替换了这些组件，允许以低得多的计算开销可微地学习软功能子结构。

#### 图神经网络的可解释性
解释 GNN 预测催生了一系列事后解释方法。GNNExplainer (Ying et al., 2019) 通过学习连续边和特征掩码，最大化选定子图与模型预测之间的互信息，提供实例特定的子结构解释。Integrated Gradients (Sundararajan et al., 2017) 提供了互补的归因框架，通过沿从基线输入到观测示例的路径积分梯度。这些方法已成为探测图模型的标准工具，评估通常基于忠实度、稀疏性和表征等指标（Pope et al., 2019; Yuan et al., 2022）。在这项工作中，我们采用这两种方法来评估在冻结的 PLM 嵌入上进行结构感知推理是否会产生具有生物学意义且可审计的解释。

## 3 问题与评估标准

我们旨在学习一个函数分类器 $f_\theta$，作用于蛋白质接触图，为其预测产生稀疏的、生物学忠实的解释。解释是一对关于边和节点特征的连续掩码 $(M, F)$。我们从两个轴评估解释：

#### 预测忠实度
应用于图解释的标准基于忠实度的指标：稀疏性、充分性（Fid+）、必要性（Fid-）和类内特征掩码稳定性（定义见附录 C）。

#### 生物学忠实度
预测指标是必要但不充分的：一个解释可能对模型忠实，但仍然没有生物学意义。我们额外要求解释在四个轴（B1）-(B4）上与已建立的酶生化一致：

#### （B1）催化残基富集
对于每个 EC 类 $c$ 和氨基酸 $a \in \Sigma$，令 $\hat{p}_{c,a}^{\mathrm{top}}$ 和 $\hat{p}_{c,a}^{\mathrm{bg}}$ 分别为 $a$ 在 $\mathcal{I}_{0.20}(G)$ 和 $V$ 中的经验频率。对数富集为 $\mathrm{Enr}(c,a) = \log_2 \frac{\hat{p}_{c,a}^{\mathrm{top}}+\delta}{\hat{p}_{c,a}^{\mathrm{bg}}+\delta}, \quad \delta=10^{-6}$. 我们测试假设 $\mathbb{E}_a[\mathrm{Enr}(c,a) \mid a \in \mathrm{Cat}] > \mathbb{E}_a[\mathrm{Enr}(c,a) \mid a \notin \mathrm{Cat}]$，其中 $\mathrm{Cat}=\{\mathrm{H,C,S,D,E,K,R,Y}\}$.

#### （B2）活性位点埋藏
重要和非重要残基之间的预期 SASA 差距：
$$ \Delta_{\mathrm{SASA}}(c) = \mathbb{E}_{i \in \mathcal{I}_{0.20}(G)}[\mathrm{SASA}_i] - \mathbb{E}_{i \notin \mathcal{I}_{0.20}(G)}[\mathrm{SASA}_i] < 0. $$

#### （B3）空间共定位
令 $\bar{D}(I) = (\|I\|_2)^{-1} \sum_{i,j \in I, \|i-j\|>20} \|i-j\|$。重要集的分布应超过非重要集，并且 C$\alpha$–C$\alpha$ 距离分布应在催化三联体区域 $d_{ij} \in [6,10]$ Å 处达到峰值。

论文的其余部分介绍了在不同结构表达水平上参数化 $f_\theta$ 的架构（第 4 节），$\mathcal{L}_{\mathrm{cls}}$ 目标的结果（第 5 节），以及同时针对忠实度（第 6 节）和生物学标准（B1）-（B4）的联合验证，我们在第 6.2 节中报告所有这些，旨在获得一个同时**准确**且**可审计**的模型。

## 4 方法

**图 1: 概述**
(a) PLM-到图投影
(b) 可微 Blob 分区
(c) 可解释的结构输出

1.  **ESM-2 PLM 密集嵌入 & 3D 坐标**
    蛋白质序列 $S \in \Sigma^N$
    ESM-2 密集嵌入 $\mathbf{X} \in \mathbb{R}^{1318}$
    3D 坐标 $C \in \mathbb{R}^{N \times 3}$

2.  **SoftBlobGIN 可微 GS**
    半径图 $\varepsilon=8$ Å
    GINEConv 骨干
    节点特征（1318-d）
    $x_i = [\, \mathbf{\phi^{\mathrm{esm}}} \, \| \, \phi^{\mathrm{phys}} \, \| \, \phi^{\mathrm{sasa}} \, \| \, \dots \,]$
    GINE 消息传递
    $h_i^{(\ell)} = \mathrm{MLP}\Big( (1+\epsilon)h_i^{(\ell-1)} + \sum_{j \in \mathcal{N}(i)} \dots \Big)$
    Gumbel-Softmax 分配
    $A_{ik} = \frac{\exp((L_{ik}+g_{ik})/\tau_t)}{\sum_{k'=1}^K \exp((L_{ik'}+g_{ik'})/\tau_t)}$
    学习到的 Blob 分区
    $b_k = g_\psi\! \left( \mathrm{LN}\! \left( \frac{\sum_i A_{ik} h_i^{(L)}}{\sum_i A_{ik} + \epsilon} \right) \right)$
    双图嵌入
    $z(G) = \big[ \, \max_k b_k \, \| \, \frac{1}{N}\sum_i h_i^{(L)} \big]$

3.  **可解释结构 & EC**
    $E, \mathbf{E}_{\mathrm{attr}}$
    $H^{(L)}$
    $\{b_k\}_{k=1}^K$

    **解释目标**
    $\min_{M,F} -\log[f_\theta(G_M)]_{\hat{y}} + \lambda_1 \|M\|_1 \dots$
    *(恢复活性位点 & 催化三联体)*