ProtSAE: Disentangling and Interpreting Protein Language Models via Semantically-Guided Sparse Autoencoders¶
会议: AAAI 2026
arXiv: 2509.05309
代码: nju-websoft/ProtSAE
领域: 生物信息学
关键词: 蛋白质语言模型, 稀疏自编码器, 语义引导, 可解释性, 特征解纠缠, 本体嵌入
一句话总结¶
提出 ProtSAE，在稀疏自编码器训练中引入语义标注和领域本体知识作为引导信号，解决传统 SAE 的语义纠缠问题，使蛋白质语言模型的隐层特征与生物学概念（分子功能、生物过程、离子结合位点等）精准对齐，同时保持高重建保真度并支持概念级别的生成控制。
研究背景与动机¶
-
蛋白质语言模型快速发展但内部机制不透明：ESM2 等 PLM 已广泛用于功能预测、结构建模和蛋白质设计，但其隐层特征如何映射到结合口袋、翻译后修饰、折叠家族等生物概念仍然未知。
-
稀疏自编码器是理解 LLM 内部表示的有力工具：SAE 基于线性特征叠加假设，将高维隐层表示分解为稀疏特征，再通过标注数据做相关性分析来解释每个特征的含义，已成功应用于 Claude 3 Sonnet 和 GPT-4 的可解释性研究。
-
传统 SAE 存在严重的语义纠缠问题：单个神经元往往同时混合了多个语义不相关的生物概念（如同一个激活同时关联 iron ion binding 和 sodium transport），导致特征解释模糊、生成控制不可靠。
-
现有方法仅在训练后做标注解释，无法从根源上解决纠缠：标准流程是先无监督训练 SAE，再用标注数据事后分析特征与概念的相关性。这种后验方式无法约束训练过程，纠缠问题得不到缓解。
-
蛋白质领域拥有丰富的结构化先验知识：Gene Ontology 等本体定义了概念间的 is-a、part-of、regulates 等逻辑关系，这些专家知识可以为特征学习提供额外的语义约束，但此前未被利用。
-
因此需要一种在训练阶段就融入语义引导的 SAE 方法：本文提出 ProtSAE，将标注数据和本体知识同时引入 SAE 训练，从源头约束特征-概念的对应关系，实现真正的语义解纠缠。
方法详解¶
整体框架¶
ProtSAE 基于 TopK-SAE 架构，将隐层激活分为两组：\(m\) 个「定义激活」\(\mathbf{z}_{\text{def}}\) 与预定义生物概念一一绑定，其余 \(n-m\) 个「未知激活」\(\mathbf{z}_{\text{unk}}\) 自由捕获未知语义。训练目标由三部分组成：重建损失 \(\mathcal{L}_{\text{rec}}\)、语义标注损失 \(\mathcal{L}_{\text{annot}}\)、本体公理损失 \(\mathcal{L}_{\text{axiom}}\)。通过强制激活和特征重缩放机制，确保定义激活有效参与重建。
关键设计 1：基于标注数据的语义解纠缠¶
- 做什么：引入一个概念预测器 \(\pi_{\text{pred}} = \sigma(\mathbf{W}_{\text{pred}}(\mathbf{x} - \mathbf{b}_{\text{dec}}) + \mathbf{b}_{\text{pred}})\)，用二元交叉熵损失 \(\mathcal{L}_{\text{annot}}\) 在标注数据上训练，使每个定义神经元学会检测特定生物概念的存在。
- 核心思路：通过权重绑定 \(\mathbf{W}_{\text{def}} = \mathbf{W}_{\text{pred}}^{\text{detach}} \cdot \exp(\mathbf{r}_{\text{pred}})\)，让编码器的定义权重与预测器共享语义方向，但通过可学习缩放因子 \(\mathbf{r}_{\text{pred}}\) 独立调整幅度。detach 操作阻止重建梯度污染预测器的语义方向。
- 设计动机：如果让重建损失直接更新 \(\mathbf{W}_{\text{pred}}\)，特征方向会被重建目标扭曲，导致语义偏移。detach + 指数缩放的设计在保持语义纯净的同时允许幅度适配重建需求。
关键设计 2：强制激活（Forced Activation）¶
- 做什么：当概念预测器判定某概念存在（\(\pi_{\text{pred}} > 0.5\)）但对应的定义激活值低于 \(\mathbf{z}_{\text{unk}}\) 的均值时，用语义偏置 \(\mathbf{z}_{\text{bias}}\) 将激活值拉升到均值以上。
- 核心思路：\(\mathbf{z}_{\text{bias}} = \mathbb{1}_{\pi_{\text{pred}}>0.5} \cdot \text{ReLU}(\text{mean}(\mathbf{z}_{\text{unk}}) - \hat{\mathbf{z}}_{\text{def}})\)，保证被预测为活跃的语义特征在解码时不会被忽略。
- 设计动机：实验观察到重建过程更依赖纠缠的无监督激活 \(\mathbf{z}_{\text{unk}}\)，定义激活容易被边缘化。强制激活确保语义特征实际参与重建，为后续的生成控制（steering）奠定基础。
关键设计 3：基于领域知识的本体约束¶
- 做什么：利用 ELEmbeddings 方法，将 Gene Ontology 中的四类标准化公理（子类 NF1、合取子类 NF2、存在性包含 NF3、存在性限制 NF4）编码为预测器权重空间中的几何约束，总损失 \(\mathcal{L}_{\text{axiom}} = \mathcal{L}_{\text{NF1}} + \mathcal{L}_{\text{NF2}} + \mathcal{L}_{\text{NF3}} + \mathcal{L}_{\text{NF4}}\)。
- 核心思路：论文证明概念预测器的权重 \(\mathbf{W}_{\text{pred}}\) 在结构上等价于 ELEmbeddings 中的本体嵌入，因此可以直接在该空间上施加本体公理约束，无需额外的嵌入空间。
- 设计动机：生物概念之间不是相互独立的（如「lytic vacuole」is-a「vacuole」），仅靠标注数据无法完整捕获这些层次关系。本体约束使特征空间的几何结构忠实反映概念间的逻辑关系，提升语义一致性。
关键设计 4：训练策略¶
- 做什么：总损失 \(\mathcal{L} = \|\hat{\mathbf{x}} - \mathbf{x}\|_2^2 + \lambda_{\text{annot}} \mathcal{L}_{\text{annot}} + \lambda_{\text{axiom}} \mathcal{L}_{\text{axiom}}\)，其中 \(\lambda_{\text{annot}} = \lambda_{\text{axiom}} = 1\)。
- 核心思路：三项损失各司其职——重建损失保持保真度，标注损失引导语义对齐，公理损失建模概念间关系。TopK 稀疏性由 \(K \in \{50, 100, 500, 1000\}\) 控制。
- 设计动机：等权组合简化了超参数调节，且实验表明这一设置在保真度和可解释性之间取得了良好平衡。
实验关键数据¶
表 1：基于探针的蛋白质功能预测（三个本体平均）¶
| 方法 | \(F_{\max}\uparrow\) | \(S_{\min}\downarrow\) | AUPR \(\uparrow\) | AUC \(\uparrow\) |
|---|---|---|---|---|
| SpLiCE | .417 | 23.4 | .360 | .329 |
| Naive SAE | .421 | 23.3 | .340 | .511 |
| Gated SAE | .441 | 22.7 | .368 | .533 |
| TopK SAE | .444 | 22.7 | .379 | .565 |
| Linear Probe (PLM) | .537 | 20.9 | .522 | .751 |
| ProtSAE | .579 | 20.9 | .487 | .797 |
ProtSAE 在所有指标上大幅超越 SAE 基线，AUC 从 TopK SAE 的 0.565 跃升至 0.797（+41%），甚至超越直接在 PLM 隐层上做线性探针的结果（AUC 0.751）。
表 2：消融实验关键对比（BPO 数据集，K=100 时的趋势）¶
| 变体 | AUC 变化 | 重建保真度变化 | 定义激活参与率 |
|---|---|---|---|
| 完整 ProtSAE | 基准 | 基准 | ~100% |
| 去掉 detach | AUC 剧降 | 略有提升 | — |
| 去掉 \(\mathcal{L}_{\text{axiom}}\) | AUC 明显下降 | 明显下降 | — |
| 去掉 \(\mathbf{z}_{\text{bias}}\) | AUC 略降 | 略降 | 明显下降 |
| 去掉 \(\mathbf{r}_{\text{pred}}\) | AUC 下降 | 下降 | — |
消融实验证明每个组件都不可或缺：detach 是保证语义纯净的关键，本体约束对复杂概念关系建模至关重要，强制激活确保定义特征参与重建。
关键发现¶
-
语义引导训练显著提升了特征-概念对齐质量：ProtSAE 的 relevance-based F1 评估中，top-10 特征的均值和最大值均远超所有 SAE 基线，说明训练时引入标注约束比事后标注更有效。
-
ProtSAE 在不同稀疏度下都保持优势：随着 \(K\) 从 50 到 1000 变化，ProtSAE 的 AUC 始终领先其他方法，同时 Loss Recovered 与 TopK SAE 相当，表明语义引导不以牺牲重建质量为代价。
-
学到的特征能可视化对应到蛋白质结构上的功能区域：铁离子结合特征精准激活在 TonB 依赖受体区域，钠离子转运特征定位到跨膜α螺旋段，金属离子结合特征能标记出具体的结合位点。
-
概念级生成控制实验表明特征具有因果意义：对特定概念激活做干预后，生成蛋白的 TM-score 显著提升、RMSD 显著下降，且 pLDDT 评分提高，说明ProtSAE 的定义特征能有效引导 PLM 生成具有目标功能的蛋白质。
-
生成的蛋白与天然蛋白结构高度相似但序列新颖：如 DNA 转录抑制因子概念干预后生成的蛋白与 A0A346G484 的 TM-score 达 0.829，但序列同一性仅 30%，体现了真正的功能导向设计能力。
亮点与洞察¶
- 将可解释性工具从「事后解释」升级为「训练引导」：这是本文最核心的思想转变——不再被动地分析 SAE 学到了什么，而是主动告诉 SAE 应该学什么，从根源上解决纠缠问题。
- detach + 缩放的权重绑定设计精巧：既保证编码器和预测器共享语义方向，又通过梯度隔离避免重建目标对语义的干扰，体现了对多任务学习中梯度冲突的深刻理解。
- 本体知识的无缝集成：通过证明预测器权重等价于 ELEmbeddings，省去了额外的嵌入空间，直接在 SAE 参数上施加公理约束，优雅且高效。
- 跨领域方法论价值：ProtSAE 的语义引导思想不仅适用于蛋白质，对任何有标注本体的 LLM 可解释性研究都有潜在的迁移价值（如医学本体、法律分类体系等）。
局限性/可改进方向¶
- 概念覆盖受限于标注数据：\(m\) 个定义激活依赖已有的 GO 标注，对于未标注的新功能或罕见概念无法建模。可考虑半监督或主动学习方式扩展概念覆盖。
- 仅在 ESM2-15B 上验证：未测试对其他 PLM（如 ProtTrans、Ankh）的适用性，泛化性存疑。
- 计算开销较大：在 ESM2-15B 上训练 SAE 需要 4 张 A800，激活宽度最大 40,000，大规模应用时成本不可忽视。
- 概念粒度固定：当前概念粒度由 GO 术语决定，缺乏层次化或多粒度的灵活解释能力。
- 生成控制实验规模有限：仅对 7 个概念做了 steering 实验，50% 掩码设定也较特殊，需要更大规模和多样设定的验证。
相关工作与启发¶
- SAE for LLM interpretability：Anthropic (Paulo et al. 2024) 和 OpenAI (Gao et al. 2024) 分别在 Claude 和 GPT-4 上验证了 SAE 的可扩展性，ProtSAE 将其扩展到蛋白质领域并解决了语义纠缠问题。
- 蛋白质可解释性：CB-pLM (Ismail et al. 2025) 用 concept bottleneck 层实现可控生成，思路互补——CB-pLM 改模型架构，ProtSAE 改解释工具。
- 本体嵌入 ELEmbeddings：Kulmanov et al. (2019) 的本体表示学习方法被巧妙地嵌入到 SAE 训练中，展示了将知识图谱引入深度学习可解释性的新范式。
- 启发：语义引导的思路可延伸至视觉模型的 SAE 分析（用 ImageNet 类别层次作为本体）、多模态模型的概念级 steering 等场景。
评分¶
- 新颖性: ⭐⭐⭐⭐ 将语义引导引入 SAE 训练是清晰的方法论创新，权重绑定和本体约束的设计精巧
- 实验充分度: ⭐⭐⭐⭐ 可解释性可视化、探针评估、消融实验、steering 实验全面覆盖，生物学验证充分
- 写作质量: ⭐⭐⭐⭐ 问题动机清晰，方法推导严谨，跨领域读者也能理解核心思想
- 价值: ⭐⭐⭐⭐ 对蛋白质 AI 可解释性有重要贡献，语义引导的思想具有广泛迁移潜力