Interpreting Latent Protein Language Model Features with Geometric Annotations¶

会议: NeurIPS 2026
arXiv: 2608.26419
OpenReview: https://openreview.net/forum?id=9fXSdER4Lv
代码: https://github.com/siddharthsetlur/ProteinLens
领域: 计算生物
关键词: 蛋白质语言模型, 稀疏自编码器, 机制可解释性, 几何特征注释, 残基级解析

一句话总结¶

针对蛋白质语言模型（pLM）潜空间特征难以精确解释的问题，提出利用主链 \(C_\alpha\) 局部几何特征自动标注稀疏自编码器（SAE）神经元，实现残基级结构特征解析并成功泛化至宏基因组未注释序列。

研究背景与动机¶

以 ESM-2 为代表的蛋白质语言模型通过海量未对齐序列的掩码预训练，展现出惊人的生物结构与功能表征能力，不仅催生了 ESMFold 等高精度从头折叠模型，更成为解读海量无同源多序列比对（MSA）信息的宏基因组“暗物质”序列的核心工具。然而，神经网络多义性使得单个隐藏神经元往往同时混杂编码多个不相关概念，机制可解释性因此面临巨大障碍。近年引入的稀疏自编码器（SAE）能够将密集的残基嵌入解耦为上万维稀疏过完备特征，极大推动了单义性生物语义的挖掘。

然而，现有的 SAE 特征注释流程高度依赖宏观生物数据库（如 InterPro、Pfam、CATH）或利用大语言模型总结强激活序列的文献背景。这种思路的核心局限在于：数据库与文本标签本质上是蛋白质级别的全局归纳，无法解析 SAE 特征在三维空间中高特异性聚焦的局部残基级几何模式；当多个不同神经元同时命中同一个广泛的结构域标签时，现有方法无法甄别其亚结构差异；面对文献库完全空白的未知微生物宏基因组序列，数据库和 LLM 注释更会彻底失效。

为了打破对人工先验标签的依赖，本文的核心 idea 是将蛋白质主链局部微观几何属性转化为可量化的特征向量，通过可控置信度的监督预测判定 SAE 特征的几何单义性，实现残基级、高覆盖度且可迁移至宏基因组的自动化结构解释。

方法详解¶

整体框架¶

整体分析流程围绕 ESM-2 8M 的残基流表征解耦、残基级激活二值化、主链局部几何特征抽取与梯度提升树（GBM）预测、以及严格的错误发现率（FDR）控制统计显著性检验展开。首先，在从 SwissProt 聚类抽取的序列上训练每层 32 倍过完备的 ReLU SAE；随后，针对给定 SAE 特征筛选代表性激活蛋白并提取对应的 AlphaFold 预测主链 \(C_\alpha\) 局部几何特征；接着训练 GBM 分类器预测该特征是否在特定残基上激活，并辅以置换检验与 Benjamini-Hochberg（BH）校正；最后与 InterPro 数据库基线、MEME 序列模序及序列绝对位置进行对比与消融分析。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["蛋白质序列<br/>SwissProt / 宏基因组"] --> B["ESM-2 8M 残基嵌入<br/>320维隐藏状态流"]
    B --> C["过完备 ReLU SAE 编码<br/>解耦为10,240维稀疏特征"]
    C --> D["残基激活二值化与过滤<br/>80%分位数阈值划分正负样本"]
    D --> E["Cα 主链局部几何特征抽取<br/>21残基窗口构建44维几何与理化向量"]
    E --> F["GBM 几何打分器与显著性检验<br/>残基混洗置换检验与 FDR 假发现率控制"]
    F --> G["多维对比与下游验证<br/>亚结构细分 / 宏基因组迁移 / 接触图消融"]

关键设计¶

1. 过完备 ReLU 稀疏自编码器架构：解耦残基级高密表征

针对 ESM-2 隐藏状态中严重的神经元多义性现象，在模型的中间残基流（第 2、4、6 层）上分别独立训练单隐藏层过完备 ReLU 稀疏自编码器。输入为第 \(l\) 层第 \(i\) 个残基处的 320 维激活向量 \(x^{(l)}_i\)，编码器将其线性映射扩展 32 倍至 10,240 维并经由 ReLU 激活截断产生稀疏特征向量，解码器随后进行重构并施加单位 \(L_2\) 范数约束。训练目标由重构均方误差与稀疏性正则惩罚项联合构成：

\[\mathcal{L}_{\text{SAE}}(x^{(l)}, \hat{x}^{(l)}) = \|x^{(l)} - \hat{x}^{(l)}\|_2^2 + \lambda \|f(x^{(l)})\|_1\]

通过在 50,000 条经由 MMseqs2 去冗余聚类（30% 序列相似度、80% 双向覆盖度）的 SwissProt 序列上优化，模型在第 6 层实现了高达 100% 的下游语言模型预测损失恢复率与超过 96% 的方差解释率，每残基激活特征数有效压缩至数百个以内，为后续下游单义性分析奠定了高保真、低重叠的特征基底。

2. 44 维多尺度主链 \(C_\alpha\) 局部几何描述符构建：从三维空间捕获构象指纹

为检验 SAE 特征是否在微观层面上响应特定几何模序，对于每个激活残基以半窗宽 \(h=10\)（即 21 残基窗口）提取 AlphaFold 预测结构的主链 \(C_\alpha\) 坐标点集，计算 44 维几何与理化协变量向量 \(\phi_i(s)\)。基础微积分几何量涵盖切向量差分所定义的曲率 \(\kappa_i = \|T_{i-1} \times T_i\|\)、四点二面角反映的带符号扭转角 \(\tau_i = \text{atan2}(y, x)\)、以及窗口内协方差矩阵最小特征值与迹之比度量的局部平面度 \(\pi_i = \lambda_{\min}(C_i) / \text{tr}(C_i)\)。

为捕获不同空间尺度与拓扑约束，特征进一步统计了子窗口（N 端、中心、C 端各三分之一）和变跨度窗口（\(h/2\) 与 \(2h\)）下的曲率、扭转与平面度的均值及标准差，辅以端到端距离比 \(E\) 与表征切线平直性的主链切向一致度 \(\nu_i\)。此外，特征还融合了排除本局部窗口后的外部空间接触统计（8Å 和 12Å 截断下的接触密度及序列距离超 12 残基的远端接触数）与 5 类基本氨基酸粗粒化理化组分比例（Gly/Pro 柔性基团、疏水、带电、极性、芳香族），全面刻画局部刚性、折叠堆积与立体构象。

3. 基于 GBM 的连续注释评分与残基混洗 FDR 显著性检验：消除模型过拟合伪影

为自动化评估几何特征对特定 SAE 神经元激活的解释力，针对每个 SAE 特征收集前 500 个最高激活蛋白质，将各蛋白内部非零激活的 80% 分位数设为正类阈值 \(\gamma\)。利用抽取的 44 维几何特征向量训练轻量级梯度提升决策树（GBM，80 棵树，最大深度 3），直接预测残基激活标签并输出连续几何预测概率分 \(\hat{y}^{\text{geom}}_{r,i}\)，采用 PR-AUC 进行无偏度量。

由于神经网络与树模型极易在不平衡稀疏数据上拟合虚假关联，本文引入了严格的蛋白质内残基混洗置换检验：保持每个蛋白质内部正负残基的总数不变，随机置乱激活位置以完全破坏其与主链几何的三维对齐，进行 100 次置换生成经验零假设分布并计算置换 \(p\) 值。在对 10,240 个特征的大规模检验中，采用 Benjamini-Hochberg（BH）过程独立对各层和各方法控制假发现率（FDR），将校正后的 \(q < 0.05\) 作为统计显著注释的严格阈值。

损失函数 / 训练策略¶

SAE 训练采用 Adam 优化器，学习率在 \(5\times 10^{-5}\) 至 \(1\times 10^{-3}\) 区间内进行对数均匀采样贝叶斯超参数优化，设置 40,000 步左右的线性预热与退火阶段，总训练步长为 500,000 步。解码器权重在每次梯度反向传播后显式投影至正交补空间并重新归一化至单位球面上。GBM 训练采用收缩率 0.1、子采样比例 0.8，叶节点最小样本量自适应设为 \(\max(5, \lfloor 0.02 \times N \rfloor)\)，严格防止树结构过度记忆少数孤立样本。

实验关键数据¶

主实验¶

在 ESM-2 8M 的第 2、4、6 层上，全面比对不同注释管道在严格 FDR 控制（\(q < 0.05\)）下的特征覆盖率，同时列出高置信度几何特征（PR-AUC 分布）的表现。

网络层	全方法并集覆盖率 (%)	蛋白质级 InterPro (%)	残基级 InterPro (%)	序列位置基线 (%)	序列模序 MEME (%)	本文几何方法覆盖率 (%)
Layer 2	86.53	56.11	42.03	62.22	69.33	84.38
Layer 4	93.50	56.75	61.46	79.19	77.56	92.79
Layer 6	94.57	56.68	67.70	71.35	77.20	94.49

在 10,000 条完全未参与训练的独立测试蛋白质上评估，几何注释显示出扎实的预测性能（PR-AUC 远超随机基线 0.01）：在 Layer 4 中，有 18.29% 的特征达到中高预测强度（PR-AUC > 0.3），其中 6.05% 的特征达到强预测（PR-AUC > 0.6）。

消融实验：亚结构细分与宏基因组迁移¶

为验证几何特征是否真正突破了粗粒度生物标签的局限，评估其对共享相同残基级 InterPro 标签特征的分离度，以及向缺乏任何数据库注释的 NMPFams 宏基因组蛋白家族（共 50,000 个家族，逾 1,000 万条序列）的零样本迁移表现。

检验场景 / 网络层	评估指标	观测数值	对照基线 / 统计说明
Layer 2: 共享 InterPro 标签细分	几何可区分特征比例 (Cosine < 0.5)	59.17%	评估 120 个 InterPro 家族，中位余弦相似度显著分离
Layer 4: 共享 InterPro 标签细分	几何可区分特征比例 (Cosine < 0.5)	77.51%	评估 578 个 InterPro 家族，达高分辨率亚结构拆解
Layer 6: 共享 InterPro 标签细分	几何可区分特征比例 (Cosine < 0.5)	68.89%	评估 45 个 InterPro 家族，有效揭示结构域内部拓扑
Layer 4: 宏基因组家族零样本迁移	匹配宏基因组家族数 (PR-AUC > 0.5)	3,875 (7.75%)	覆盖 757,802 条无注释宏基因组序列，中位 PR-AUC 达标
Layer 6: 宏基因组家族零样本迁移	匹配宏基因组家族数 (PR-AUC > 0.5)	8,108 (16.22%)	覆盖 1,588,446 条无注释宏基因组序列，展示强大泛化力

关键发现¶

超越序列与宏观数据库：几何注释在各层单独覆盖率均超过 84%~94%，远高于残基级 InterPro 的 42%~67%。大量被 InterPro 标记为空白的 SAE 特征被证明具有高度确定的局部几何构象（如 \(\beta\)-折叠扭曲转角、接触紧凑度）。
细分粗粒度生物学标签：以蛋白 Q15120 为例，四个同时被标记为“组氨酸激酶/HSP90 样 ATPase 结构域”的 SAE 特征（f/6775, f/5508, f/8254, f/9608）被几何特征精准分离为该结构域前段螺旋、中段螺旋、连接区与尾端螺旋，特征间残基重叠率极低（中位余弦 0.25）。
干预消融呈现因果响应：在 Layer 4 对特征 f/373（主导特征为 12Å 序列距离）进行连续消融实验，随着消融强度增加，ESM-2 内部预测的接触图在序列间距 [13, 48] 范围内的接触概率发生线性单调偏移，证明几何特征非单纯相关而是具备因果调制效应。

亮点与洞察¶

将纯序列语言模型的潜变量对齐到物理空间微观几何：不同于前人利用 100 残基 RMSD 做粗略刚体比对，本文构建了 44 维严密的主链微分几何描述符，使得解释深入到单残基曲率、扭转与局部接触层级。
开箱即用的宏基因组“暗物质”功能与结构挖掘：将 SwissProt 上训练的几何树模型直接套用于缺乏任何生物先验注释的宏基因组蛋白，成功为超过 150 万条微生物序列指认了微观局部结构模序。
开源可视化平台 GeoPedia：开发了网页端三维分子结构与特征激活联动交互平台，为结构生物学家与机制可解释性研究者架设了直观的交互桥梁。

局限与展望¶

作者承认的局限：受制于算力，消融实验主要在缺乏完整折叠头（Folding Head）的 ESM-2 8M 上开展，干预结果仅限于接触图层级；此外几何特征提取依赖 AlphaFold 预测结构，在低 pLDDT 区域或特殊极端构象下可能存在潜在预测偏差。
深度审视局限：GBM 模型虽然解释性好，但特征工程本质上固定在 \(h=10\) 的固定窗口内，对超长程非局部协同变异的敏感度有限；同时置换检验主要固定了训练后的树参数，更严格的重训练置换虽然在 Layer 4 抽样显示仅减少约 0.3% 的特征，但并未全量铺开。
改进与拓展方向：未来可将几何注释管线拓展至 ESMFold 完整端到端模型，通过特征导向（Feature Steering）直接干预三维骨架生成，赋能高活性人工嵌合酶设计与定向进化断点选择。

评分¶

新颖性: ⭐⭐⭐⭐⭐ 首次将系统性三维主链微分几何描述符与稀疏自编码器结合，实现残基级机制可解释性。
实验充分度: ⭐⭐⭐⭐⭐ 涵盖 6 层 SAE、50k 序列、FDR 严格统计检验、宏基因组迁移验证及因果消融分析。
写作质量: ⭐⭐⭐⭐⭐ 逻辑结构严谨，统计学定义清晰，跨学科叙事连贯。
价值: ⭐⭐⭐⭐⭐ 为蛋白质基础模型开辟了物理与几何维度的白盒剖析途径，对结构生物学和蛋白质工程具有高度实用价值。