问题:PLM 表征的黑盒困境
蛋白质语言模型(PLM)如 ESM2 通过掩码语言建模在数亿条蛋白质序列上无监督预训练,产生的表征在结构预测、功能注释等下游任务中表现卓越。但这些表征是黑盒——我们知道 480 维向量编码了某种生物学信息,却不知道每一维对应什么生物学概念,也不知道为什么特定表征导致特定预测。
这种不可解释性带来三重挑战:
- 人机协作障碍:研究者无法理解模型为何做出特定预测,难以信任和调试
- 生物学洞察受限:表征中可能蕴含未知的生物学规律,但无法系统提取
- 安全风险:开源 PLM 可能无意中编码可用于生物武器的敏感信息,需要审计手段
已有研究分析了 PLM 的注意力头和内部嵌入如何捕获接触图和结合位点,但这些方法不提供无监督地系统提取特定表征中生物学有意义特征的途径。本文的核心贡献正是填补这一空白。
SAE 方法论:从超叠加到单语义特征
超叠加问题
神经网络中的神经元是多语义的(polysemantic):一个神经元会被多个不相关的特征激活。原因在于——真实世界中特征出现稀疏,模型需要表达的特征数远多于层中神经元数。因此,网络将多个稀疏特征超叠加(superposition)存储在较少的神经元中,而非为每个特征分配一个专用神经元。这导致网络性能良好但难以解释。
在 NLP 中,这一问题已被广泛研究。稀疏自编码器(SAE)被证明可以有效提取语言模型中的单语义特征——在 NLP 中发现了如希伯来文字、DNA 序列、数学符号等可解释特征。
SAE 架构:TopK 稀疏自编码器
本文使用 TopK SAE(Gao et al. 2024),这是解释 LLM 最流行的 SAE 变体之一。核心思想:自编码器的隐藏层比输入更宽(而非更窄),并通过稀疏约束使隐藏层神经元稀疏激活,从而将多语义神经元分解为单语义的稀疏特征。
具体架构:输入 x_L ∈ ℝ⁴⁸⁰(ESM2 第 L 层表征),编码器为:
其中 TopK 函数只允许 top-k 个值通过,其余置零。y ∈ ℝ²⁰⁰⁰⁰ 是稀疏特征向量。解码器为:
训练目标为标准 MSE 重建损失。为解决大量神经元死亡问题,引入辅助重建项(aux reconstruction)。
训练设置
| 参数 | 蛋白质级 SAE | 氨基酸级 SAE |
|---|---|---|
| 基础模型 | ESM2 t12_35M_UR50D | 同左 |
| 训练数据 | 5000 万序列(UniRef50,≤1024 残基) | 同左 |
| 输入维度 | 480(mean-pooled 蛋白质表征) | 480(单个氨基酸表征) |
| 隐藏层维度 | 20,000 | 20,000 |
| TopK 稀疏度 k | 64 | 16 |
| 训练层 | 第 6-10 层(共 12 层) | 第 6-10 层 |
| 训练/验证划分 | 80/20 | 80/20 |
选择第 6-10 层而非最后两层,是因为借鉴 LLM 可解释性经验:最有趣的解释层在模型末尾附近但非最末端——最末端层通常专门用于无监督训练任务(token 预测)。
蛋白质级 vs 氨基酸级表征
论文在两个粒度上训练 SAE:
- 蛋白质级:对序列维度 mean-pooling 得到单一 480 维向量,这是下游任务最常用的表征形式。解释这些表征有助于理解和调试下游预测
- 氨基酸级:直接对每个氨基酸的 480 维 token 表征训练 SAE,揭示 ESM2 内部跟踪的特征。为保持可解释性协议可行性,对 SAE 特征在序列上 mean-pooling 后解释
转码器:学习层间转换的稀疏近似
论文引入转码器(Transcoder)——SAE 的变体,源自 NLP 中 Dunefsky et al. 的工作。与 SAE 的关键区别:
| 属性 | SAE | Transcoder |
|---|---|---|
| 输入 | 第 L 层表征 x_L | 第 L 层表征 x_L |
| 目标输出 | 重建 x_L 自身 | 预测第 L+1 层表征的增量 x_{L+1} − x_L |
| 学习内容 | 表征的稀疏分解 | 层间计算的稀疏近似 |
| UndoLayerNorm | 是 | 否(输入输出来自不同层) |
转码器学习从一层蛋白质表征到下一层的转换中哪些稀疏特征参与。取差值 x_{L+1} − x_L 而非 x_{L+1} 是因为 ESM2 中的残差连接——转码器只需学习该层新增的部分,而非包含前层残差的总和。转码器特征的可解释性与 SAE 相当,但提供了理解 PLM 内部信息流的不同视角。
GO 富集分析:稀疏特征与功能注释的关联
为评估蛋白质级 SAE 特征的生物学可解释性,论文对 18,142 个人类蛋白质进行了 GO 富集分析。对每个稀疏特征,收集所有非零激活的蛋白质,使用超几何检验检测 GO 术语过表达,Benjamini-Hochberg 校正控制 FDR。
GO 单语义性量化
论文提出两个量化指标:
- GO 单语义性:对每个特征的富集 GO 术语计算 GO DAG 中的成对最短路径长度。路径越短 → GO 术语越接近 → 特征越"单语义"
- GO DAG 位置:计算富集 GO 术语的最低共同祖先(LCA)深度。深度越大 → 功能越特异(叶节点级);深度越小 → 功能越通用(根级)
关键发现:跨 GO 层级的单语义特征
在第 7 层 SAE 中,大量特征在所有 GO DAG 深度上都表现出低最短路径长度——即功能单语义性。与随机基线相比,在每个 LCA 深度区间,真实特征的最短路径长度均显著更小(p < 0.001)。这意味着 SAE 特征映射到了 GO 层级的所有水平,从非常具体的叶级节点到更广泛的功能类别。
跨层级的生物学过程特征示例
| 神经元 | LCA 深度 | 富集 GO 术语 | 生物学含义 |
|---|---|---|---|
| 2793 | 1 | 病毒翻译移码、病毒过程、天冬氨酸型肽酶活性 | 病毒相关蛋白水解 |
| 10282 | 2 | 逆转录转座、转座 | 转座子相关蛋白 |
| 9095 | 3 | 肽修饰、肽代谢过程 | 肽链加工 |
| 16984 | 4 | 通过宿主 ESCRT 复合体正向调控病毒出芽 | 病毒包膜形成 |
| 7631 | 5 | 金属硫簇组装、铁硫簇组装 | 铁硫蛋白代谢 |
| 953 | 6 | 通过 [4Fe-4S] 簇转移的蛋白质成熟 | 铁硫簇成熟 |
| 1105 | 7 | 铁离子转运、过渡金属离子转运 | 金属离子运输 |
| 8646 | 8 | 正向调控蛋白质氧化 | 氧化调控 |
| 3823 | 9 | 白三烯 D4 分解、二肽酶活性 | 炎症介质代谢 |
| 11233 | 10 | 正向调控丝氨酸型肽酶活性 | 丝氨酸蛋白酶调控 |
| 3771 | 11 | 磷酸戊糖分流氧化支路、NADPH 再生、G6PD 活性 | 磷酸戊糖途径 |
这张表展示了 SAE 特征跨越 GO DAG 从深度 1(最通用)到深度 11(最特异)的完整覆盖——且每个特征只对应一个紧密相关的 GO 术语簇,体现了真正的功能单语义性。
GO 覆盖率
第 9 层 SAE 的富集 GO 术语总计 9,789 个(生物过程 6,024、分子功能 2,740、细胞组分 1,025),覆盖了分析蛋白质集中 52.17%、62.93%、59.83% 的 GO 术语。SAE 和转码器在同一层的富集 GO 术语高度重叠(9,180/9,789 共享),表明两种方法捕获了相似的生物学信号。
Claude 自动可解释性:LLM 解读 LLM
20,000 个稀疏特征无法人工逐一解释。论文借鉴 NLP 中的自动可解释性方法,使用 Claude 3.5 Sonnet 实现规模化特征解读。
协议流程
- 选样:从 50,000 条 UniProt 序列中提取每个特征的激活值,按激活区间分块采样——高激活区 6 条、中激活区 3 条、低激活区 2 条、非激活区 15 条。非激活区等价于随机采样(因稀疏性)
- 解读:将一半序列(8 激活 + 7 非激活)的 ID、激活值、蛋白质名称、基因名、蛋白质家族、GO 术语提供给 Claude,要求生成一句话人类可读解读
- 模拟:在新的对话中,将另一半序列的相同元数据(不含激活值)和 Claude 的解读提供给 Claude,要求预测激活值
- 评分:计算预测激活值与真实激活值的 Pearson 相关性作为可解释性分数
每个 SAE/转码器分析 200 个随机选取的充分激活特征。固定随机种子 42 确保可复现。为应对高稀疏性导致的采样不足,设计了溢出机制——当某激活区间序列不足时,从相邻区间补充。
氨基酸级特殊处理
蛋白质序列长度可变且氨基酸数远多于 NLP 中的 token 数。为保持可行性,对氨基酸级 SAE 特征在序列上 mean-pooling 后再执行上述协议——即解释每个特征在整个序列上的平均激活。这是一个局限:无法定位到特定氨基酸(如催化残基)。
蛋白质级 SAE 特征:家族与功能
Claude 自动解读揭示了蛋白质级 SAE 特征与特定蛋白质家族和功能的精确对应:
| 层 | 神经元 | Claude 解读 | 相关性 |
|---|---|---|---|
| 9 | 6384 | 磷酸烯醇式丙酮酸羧激酶(ATP)活性和糖异生 | 0.972 |
| 9 | 11982 | NAD 激酶家族,NAD 代谢过程和 NADP 生物合成 | 0.954 |
| 7 | 13811 | 二氨基庚二酸差向异构酶家族,赖氨酸生物合成 | 0.993 |
| 7 | 2579 | 甲基转移酶活性,RNA/DNA 甲基化 | 0.930 |
| 8 | 19647 | IUNH 家族,核苷水解酶活性,核碱基代谢 | 0.999 |
| 8 | 3800 | 跨膜转运离子/氨基酸,位于质膜 | 0.721 |
| 10 | 16397 | 假尿苷合成,tRNA/rRNA 假尿苷合成酶(TruB/RluA 家族) | 0.969 |
| 10 | 15683 | 电子传递链,NADH 脱氢酶和醌氧化还原酶(复合体 I) | 0.758 |
这些特征覆盖了多种生物学功能类别:
- 代谢酶:NAD 激酶(NAD 代谢)、G6PD(磷酸戊糖途径)、PEP 羧激酶(糖异生)
- 翻译机制:假尿苷合成酶(tRNA/rRNA 修饰)、PTH 家族(肽酰-tRNA 水解)
- 膜转运:跨膜离子/氨基酸转运蛋白
- 能量转导:电子传递链复合体 I 组分
- 表观遗传:甲基转移酶(RNA/DNA 甲基化)
- 核苷酸代谢:IUNH 家族核苷水解酶
氨基酸级 SAE 特征:残基级生物学信号
氨基酸级 SAE 揭示了 ESM2 在残基粒度上跟踪的生物学特征。中位 Pearson 相关性 0.673-0.714,与蛋白质级 SAE 相当。
| 层 | 神经元 | Claude 解读 | 相关性 |
|---|---|---|---|
| 10 | 7554 | 嗅觉/味觉感知,昆虫化学感受器和味觉受体 | 0.987 |
| 8 | 11965 | 磺基转移酶活性,类固醇/脂质代谢 | 0.714 |
| 9 | 15537 | PTH 家族,肽酰-tRNA 水解酶活性,翻译过程 | 0.995 |
| 10 | 16441 | 跨膜转运蛋白,主要促进因子超家族 | 0.711 |
| 9 | 19717 | UPF0312 家族(未表征),位于周质空间 | 0.981 |
| 9 | 10718 | P 型阳离子转运 ATP 酶,金属离子跨膜转运 | 0.974 |
| 7 | 10650 | 钙离子结合,肌肉相关/钙响应细胞过程 | 0.967 |
| 9 | 1346 | 钼蝶呤辅因子生物合成 | 0.990 |
| 7 | 12791 | 腺嘌呤脱氨酶活性,金属依赖水解酶超家族 | 0.998 |
氨基酸级特征的独特发现:
- 嗅觉/味觉受体(F7554, 相关性 0.987):检测昆虫化学感受器和味觉受体——一个高度特化的感觉蛋白类别
- 未表征家族(F19717, UPF0312):SAE 为功能未知的蛋白质家族提供了特征级标注,暗示 PLM 可能已编码了尚未被实验表征的生物学信息
- 钙离子结合(F10650):与肌肉相关或钙响应细胞过程关联——残基级特征捕获了功能位点的上下文
- 辅因子生物合成(F1346):钼蝶呤辅因子——一个精细的代谢通路节点
转码器特征:层间计算的稀疏分解
转码器特征的可解释性与 SAE 相当(中位相关性 0.66-0.73),但揭示了不同视角——哪些稀疏特征参与了从一层到下一层的信息转换。
| 层 | 神经元 | Claude 解读 | 相关性 |
|---|---|---|---|
| 9 | 9270 | 精氨酸-tRNA 连接酶,精氨酰-tRNA 氨酰化 | 0.989 |
| 9 | 12680 | 细胞凋亡,半胱天冬酶及相关蛋白 | 0.680 |
| 8 | 12186 | FMN 依赖 NADH:醌氧化还原酶,偶氮还原酶 1 型 | 0.937 |
| 8 | 12187 | 金属离子结合和转运,铁/铜等 | 0.711 |
| 8 | 15097 | 短链脱氢酶/还原酶(SDR)家族,脂肪酸延伸/视黄醇代谢 | 0.747 |
| 7 | 4745 | 趋化因子家族,免疫响应和细胞信号传导 | 0.989 |
| 10 | 19393 | UDP-糖基转移酶/UDP-葡萄糖醛酸转移酶活性 | 0.943 |
| 10 | 4745 | GTP 3',8-环化酶活性,钼辅因子生物合成(MoaA 家族) | 0.673 |
转码器特征同样对应明确的蛋白质家族和功能:tRNA 合成酶、凋亡相关半胱天冬酶、氧化还原酶、免疫趋化因子、糖基转移酶等。SAE 和转码器在第 9 层共享 9,180/9,789 个富集 GO 术语,表明两种方法捕获了高度重叠的生物学信号,但转码器额外提供了"哪些特征在层间转换中被使用"的信息。
SAE vs ESM 神经元:可解释性对比
论文系统的对比了 SAE 稀疏特征与原始 ESM2 神经元的可解释性。在所有层、所有表征类型上,SAE 特征显著优于 ESM 神经元:
| 表征类型 | 层 | SAE 中位相关性 | ESM 中位相关性 | SAE 优势 |
|---|---|---|---|---|
| 蛋白质级 | 6 | 0.706 | 0.520 | +0.186 |
| 蛋白质级 | 7 | 0.702 | 0.471 | +0.231 |
| 蛋白质级 | 8 | 0.726 | 0.454 | +0.272 |
| 蛋白质级 | 9 | 0.692 | 0.425 | +0.267 |
| 蛋白质级 | 10 | 0.756 | 0.435 | +0.321 |
| 氨基酸级 | 6 | 0.673 | 0.499 | +0.174 |
| 氨基酸级 | 9 | 0.714 | 0.416 | +0.298 |
| 氨基酸级 | 10 | 0.703 | 0.438 | +0.265 |
SAE 特征在蛋白质级和氨基酸级表征上均稳定优于 ESM 神经元约 0.17-0.32 个相关性单位。ESM 神经元中位相关性最高仅约 0.52,而 SAE 特征稳定在 0.7 附近。直方图对比显示 SAE 有大量高可解释特征(0.7-1.0),而 ESM 神经元集中在较低区间。
稀疏度统计
SAE 特征高度稀疏。以第 9 层为例:
- 蛋白质级 SAE:平均每个特征激活 0.32% 的蛋白质,中位数仅 0.02%
- 氨基酸级 SAE:平均每个特征激活 2.36% 的蛋白质,中位数 0.65%
这种极端稀疏性正是单语义性的基础——每个特征只在非常特定的蛋白质子集上激活。
序列相似性与可解释性的关系
共激活蛋白质确实比非共激活蛋白质有更高的序列相似性(全局比对、局部比对、k-mer 分数均更高),但序列相似性与可解释性分数的 Pearson 相关性较弱(0.12-0.26)。这表明可解释的 SAE 特征不仅仅是序列相似性聚类——模型捕获了超越简单序列同源性的功能信号。
SAE 与生物学特征:深度关联分析
从 GO 层级到蛋白质家族的完整映射
本文最核心的贡献是系统性地证明了 SAE 特征与生物学特征的深度关联。这种关联在多个层面体现:
1. GO 层级的全覆盖:SAE 特征映射到 GO DAG 的所有深度(1-11+),从根级通用功能(如"病毒过程")到叶级特异功能(如"磷酸戊糖分流氧化支路")。在每个深度上,特征都比随机基线显著更具单语义性(p < 0.001)。这不是偶然的——SAE 无监督地发现了与数十年生物学 curated 知识一致的功能分类。
2. 蛋白质家族的精确对应:多个 SAE 特征直接对应特定蛋白质家族:
- NAD 激酶家族(NAD 代谢 + NADP 生物合成)
- IUNH 家族(核苷水解酶)
- PTH 家族(肽酰-tRNA 水解酶)
- 二氨基庚二酸差向异构酶家族(赖氨酸生物合成)
- TruB/RluA 家族(tRNA/rRNA 假尿苷合成酶)
- SDR 家族(短链脱氢酶/还原酶)
- 趋化因子家族(免疫响应)
- UPF0312 家族(未表征——SAE 为功能未知家族提供了标注)
3. 功能维度的多样性:SAE 特征覆盖的生物学功能维度极其广泛:
- 代谢:糖异生、磷酸戊糖途径、NAD 代谢、脂肪酸延伸、视黄醇代谢
- 翻译:tRNA 氨酰化、假尿苷合成、肽酰-tRNA 水解
- 转运:跨膜离子转运、金属离子转运、氨基酸转运
- 信号传导:GTP 环化酶、趋化因子、钙离子结合
- 氧化还原:电子传递链复合体 I、NADH 脱氢酶、FMN 依赖氧化还原酶
- 表观遗传:RNA/DNA 甲基化
- 免疫:半胱天冬酶/凋亡、趋化因子
- 感觉:嗅觉/味觉受体
- 辅因子:钼蝶呤辅因子生物合成、铁硫簇组装
- 病毒:翻译移码、病毒出芽
4. 超越序列相似性的功能信号:共激活蛋白质的序列相似性与可解释性相关性弱(0.12-0.26),意味着 SAE 特征不仅捕获了序列同源性——模型学到了功能层面的抽象,能够将序列不同但功能相似的蛋白质映射到同一特征。
与 ESM Cambrian (ESMC) SAE 分析的对比
本文(Gujral et al.)与同期 ESMC 论文(Rives et al. 2026)在 SAE 应用上有重要异同:
| 维度 | Gujral et al. (PNAS 2025) | Rives et al. (bioRxiv 2026) |
|---|---|---|
| 基础模型 | ESM2 t12_35M(3500 万参数,12 层) | ESMC 6B(60 亿参数,80 层) |
| SAE 特征数 | 20,000(单一码本) | 16,384(主分析),范围 8K-131K |
| 稀疏度 | k=64(蛋白质级),k=16(AA 级) | 8-128 个激活/氨基酸 |
| 训练数据 | 5000 万序列(UniRef50) | 80 亿 token(与 LM 训练同分布) |
| 标注方法 | Claude 3.5 Sonnet(200 特征/模型) | GPT-5 多智能体五重验证(全部特征) |
| 表征粒度 | 蛋白质级 + 氨基酸级 | 仅氨基酸级(max-pooling 为蛋白质向量) |
| 特征分类 | GO 富集 + Claude 自然语言描述 | 生物学层级分类(一级→进化主题) |
| 特征层级覆盖 | GO DAG 深度 1-11 | 氨基酸→二级结构→三级→功能→进化 |
| 转码器 | 是(层间转换稀疏近似) | 否 |
| 应用 | 可解释性、安全审计 | 蛋白质图谱聚类、功能发现 |
Gujral et al. 是先驱性工作——首次将 SAE 应用于 PLM,在较小模型上验证了方法论。Rives et al. 在 170 倍大的模型上扩展了这一方法,发现了更丰富的生物学层级(从氨基酸到进化主题),并应用于 68 亿序列的图谱构建。两者互补:Gujral et al. 建立了 SAE-for-PLM 的方法论框架(包括转码器创新和 Claude 自动解读),Rives et al. 将其推向规模极限。
技术启示、局限与同期工作
核心启示
- SAE 有效解耦 PLM 黑盒:20,000 个稀疏特征在所有层和表征类型上比 ESM2 神经元更具可解释性(中位 0.7 vs 0.5),且完全无监督
- GO 富集验证生物学意义:SAE 特征与 GO 术语在所有层级强关联,从通用功能到特异通路,且比随机基线显著更具单语义性
- 转码器提供互补视角:学习层间转换的稀疏近似,可解释性与 SAE 相当,共享 94% 的富集 GO 术语
- LLM 可规模化解读 PLM:Claude 3.5 自动可解释性协议可扩展到任意数量的特征,为 PLM 审计提供工具
- 未表征家族获得标注:UPF0312 等功能未知家族被 SAE 特征捕获,暗示 PLM 已编码实验尚未发现的功能信息
- 安全审计潜力:无监督审计 PLM 是否编码危险生物信息(如生物武器相关),对开源模型治理至关重要
局限
- 模型规模:仅在 ESM2 35M 上验证,更大 PLM(如 ESM2 15B 或 ESMC 6B)的 SAE 行为可能不同
- Mean-pooling 限制:蛋白质级表征使用 mean-pooling,其他提取方式(如 BOS token)可能产生不同特征集
- AA 级 mean-pooling:氨基酸级 SAE 特征在序列上 mean-pooling 后解释,无法定位到特定残基(如催化位点)
- 仅分析 200 个特征:因 API 限制,每个 SAE 仅解读 200/20,000 个特征,更全面分析可能发现更多生物学特征
- SAE 非功能预测工具:作为无监督方法,SAE 在功能预测上无法与监督方法竞争
- TopK SAE 种子敏感性:NLP 中已记录 TopK SAE 特征可能因随机种子不同而变化,本文未评估此效应
- SAE 架构选择:仅使用 TopK SAE,L1 SAE、Gated SAE 或 BatchTopK SAE 可能更适合特定表征类型
同期工作
- Simon & Zou (Nature Methods 2025, InterPLM):在 PLM 氨基酸级表征上训练 SAE,使用 L1 SAE 和 max-pooling,同样使用 LLM 辅助可解释性
- Adams et al. (2025, InterProt):在 PLM 氨基酸级表征上训练 SAE,使用 TopK SAE 和 mean-pooling,聚焦于特征评估
- Brixi et al. (2025):在基因组基础模型 Evo2 上训练和解释 SAE
- 本文独特贡献:同时训练和解释蛋白质级和氨基酸级 SAE + 转码器,蛋白质级表征是下游任务最常用的形式
论文信息
标题:Sparse autoencoders uncover biologically interpretable features in protein language model representations
作者:Onkar Gujral, Mihir Bafna, Eric Alm, Bonnie Berger
机构:MIT Department of Mathematics / CSAIL / Center for Microbiome Informatics and Therapeutics
发表:PNAS, 2025, Vol. 122 No. 34, e2506316122
DOI:10.1073/pnas.2506316122
基础模型:ESM2 t12_35M_UR50D(3500 万参数,12 层,480 维)
SAE:TopK SAE,20,000 特征,k=64(蛋白质级)/ k=16(AA 级),第 6-10 层
转码器:层间转换稀疏近似,同架构,目标为 x_{L+1} − x_L
自动可解释性:Claude 3.5 Sonnet,200 特征/模型,Pearson 相关性评分
GO 分析:18,142 人类蛋白质,9,789 富集 GO 术语(第 9 层 SAE)
代码:https://github.com/onkarsg10/Rep_SAEs_PLMs
关键词:可解释性, 蛋白质语言模型, 稀疏自编码器, 转码器, 超叠加, 单语义, Gene Ontology, 自动可解释性