AiPaper
论文状态:已完成

InterPLM: Discovering Interpretable Features in Protein Language Models via Sparse Autoencoders

发表:2024/11/13
原文链接PDF 下载
价格:0.10
价格:0.10
已有 4 人读过
本分析由 AI 生成,可能不完全准确,请以原文为准。

TL;DR 精炼摘要

本文提出了一种利用稀疏自编码器从蛋白质语言模型(PLM)中提取可解释特征的新方法。通过分析ESM-2模型的嵌入,发现多达2548个人类可解释的潜在特征,这些特征与143个生物学概念高度相关。研究表明,PLM能够编码丰富的生物学信息,且可应用于蛋白质数据库补全和序列生成。

摘要

Protein language models (PLMs) have demonstrated remarkable success in protein modeling and design, yet their internal mechanisms for predicting structure and function remain poorly understood. Here we present a systematic approach to extract and analyze interpretable features from PLMs using sparse autoencoders (SAEs). By training SAEs on embeddings from the PLM ESM-2, we identify up to 2,548 human-interpretable latent features per layer that strongly correlate with up to 143 known biological concepts such as binding sites, structural motifs, and functional domains. In contrast, examining individual neurons in ESM-2 reveals up to 46 neurons per layer with clear conceptual alignment across 15 known concepts, suggesting that PLMs represent most concepts in superposition. Beyond capturing known annotations, we show that ESM-2 learns coherent concepts that do not map onto existing annotations and propose a pipeline using language models to automatically interpret novel latent features learned by the SAEs. As practical applications, we demonstrate how these latent features can fill in missing annotations in protein databases and enable targeted steering of protein sequence generation. Our results demonstrate that PLMs encode rich, interpretable representations of protein biology and we propose a systematic framework to extract and analyze these latent features. In the process, we recover both known biology and potentially new protein motifs. As community resources, we introduce InterPLM (this http URL), an interactive visualization platform for exploring and analyzing learned PLM features, and release code for training and analysis at this http URL.

思维导图

论文精读

中文精读

1. 论文基本信息

1.1. 标题

InterPLM: 通过稀疏自编码器发现蛋白质语言模型中的可解释特征 (InterPLM: Discovering Interpretable Features in Protein Language Models via Sparse Autoencoders)

1.2. 作者

Elana Simon 和 James Zou

1.3. 发表期刊/会议

未明确指明发表期刊或会议,但已在 arXiv 预印本平台发布。

1.4. 发表年份

2024年

1.5. 摘要

蛋白质语言模型 (PLM) 在蛋白质建模和设计方面取得了显著成功,但其预测结构和功能的内部机制仍不甚明了。本文提出了一种系统性方法,利用稀疏自编码器 (SAE) 从 PLM 中提取和分析可解释特征。通过在 PLM ESM-2 的嵌入 (embeddings) 上训练 SAE,我们识别出每层多达 2,548 个人类可解释 (human-interpretable) 的潜在特征,这些特征与多达 143 个已知生物学概念(如结合位点、结构基序和功能域)高度相关。相比之下,检查 ESM-2 中的单个神经元发现每层只有多达 46 个神经元与 15 个已知概念有清晰的概念对齐,这表明 PLM 大多数概念以叠加 (superposition) 形式表示。除了捕获已知注释外,我们还展示了 ESM-2 学习到了一些无法映射到现有注释的连贯概念 (coherent concepts),并提出了一种利用语言模型 (LLM) 自动解释 SAE 学习到的新颖潜在特征的流水线。作为实际应用,我们演示了这些潜在特征如何填充蛋白质数据库中缺失的注释,并实现蛋白质序列生成 (protein sequence generation)定向引导 (targeted steering)。我们的结果表明,PLM 编码了丰富、可解释的蛋白质生物学表示,我们提出了一个系统框架来提取和分析这些潜在特征。在此过程中,我们既恢复了已知生物学,也发现了潜在的新蛋白质基序。作为社区资源,我们推出了 InterPLM (interPLM.ai) 这一交互式可视化平台,用于探索和分析 PLM 学习到的特征,并发布了训练和分析代码 (github.com/ElanaPearl/interPLM)。

1.6. 原文链接

https://arxiv.org/abs/2412.12101 PDF 链接: https://arxiv.org/pdf/2412.12101.pdf 发布状态:预印本 (Preprint)

2. 整体概括

2.1. 研究背景与动机

2.1.1. 蛋白质语言模型 (PLM) 的崛起与黑箱问题

近年来,蛋白质语言模型 (PLM),如 ESM-2,在蛋白质建模和设计任务中展现出卓越的性能。它们能够从大量蛋白质序列数据中学习到复杂的模式,并预测蛋白质的结构和功能。然而,PLM 内部的工作机制,即它们如何从序列中提取信息并做出这些预测,仍然是一个黑箱 (black box)。这种缺乏可解释性 (lack of interpretability) 限制了我们对模型行为的理解,并阻碍了其在生物学发现和药物设计等关键领域的进一步应用。

2.1.2. 现有可解释性研究的局限性

现有的 PLM 可解释性研究通常集中于分析注意力机制或单个神经元的激活模式。然而,研究发现 PLM 中的许多神经元具有多义性 (polysemantic),即一个神经元可能同时编码多个不相关的生物学概念,这使得直接从单个神经元中提取清晰、独立的生物学概念变得困难。这种现象被称为叠加 (superposition),它意味着模型内部的概念可能以线性组合的形式存储在神经元中,而不是每个神经元对应一个概念。这构成了理解 PLM 内部表示的一个重大挑战。

2.1.3. 填补空白:对 PLM 内部机制的系统性解释

为了解决 PLM 的黑箱问题和神经元叠加现象带来的挑战,本研究旨在开发一种系统性框架,从 PLM 内部提取人类可解释 (human-interpretable) 的生物学特征。作者认为,如果能够清晰地识别并解释 PLM 内部编码的这些特征,不仅能增进对 PLM 预测能力的理解,还能促进新的生物学发现,如识别未知的蛋白质基序或功能区域。

2.2. 核心贡献/主要发现

本文的主要贡献和发现可以总结如下:

  • 引入稀疏自编码器 (SAE) 进行 PLM 特征提取: 首次系统性地将 SAE 应用于 PLM 的中间嵌入 (embeddings),成功地从 ESM-2 模型中提取出大量(每层高达 2,548 个)稀疏 (sparse)可解释 (interpretable) 的潜在特征。这些特征比单个神经元更能清晰地捕获生物学概念,有效解决了神经元叠加 (superposition) 的问题。
  • 量化可解释性并发现丰富的生物学概念: 通过与 Swiss-Prot 数据库 (Swiss-Prot database) 中 143 个已知生物学概念(如结合位点、结构基序、功能域、翻译后修饰等)进行定量关联分析,证明了 SAE 提取的特征能够比 ESM-2 的原始神经元捕获显著更多更精确的生物学注释。
  • 揭示 PLM 学习到的未标记概念: 发现 ESM-2 不仅捕获了已知的生物学概念,还学习到了一些连贯概念 (coherent concepts),这些概念目前无法直接映射到现有的 Swiss-Prot 注释,表明 PLM 具有发现新生物学模式的潜力。
  • 开发基于大语言模型 (LLM) 的自动特征解释流水线: 提出并实现了一个创新的流水线,利用 Claude-3.5 Sonnet 等大语言模型 (LLMs) 自动为 SAE 学习到的新颖潜在特征生成有意义的自然语言描述,显著提高了特征解释的效率和可扩展性。
  • 展示实际应用价值:
    • 填充缺失注释: 演示了如何利用这些可解释特征的激活模式来识别蛋白质数据库中缺失的生物学注释(例如,推断出未标记的 Nudix 盒基序或肽酶域的存在)。
    • 定向引导蛋白质序列生成: 首次展示了通过激活特定的 SAE 特征,可以定向引导 (targeted steering) PLM 生成具有特定周期性模式(例如胶原蛋白样区域中的 Gly-X-X 重复)的蛋白质序列,为蛋白质设计提供了新的工具。
  • 发布社区资源: 推出了 InterPLM 交互式可视化平台 (interPLM.ai) 和相关代码 (github.com/ElanaPearl/interPLM),使研究人员能够探索和分析 PLM 学习到的特征,促进社区的进一步研究。

3. 预备知识与相关工作

3.1. 基础概念

3.1.1. 蛋白质语言模型 (PLM)

蛋白质语言模型 (Protein Language Models, PLMs) 是一类基于深度学习 (deep learning) 的模型,它们将蛋白质序列视为一种“生物语言”,将氨基酸视为“词元 (tokens)”。通过在海量的蛋白质序列数据集上进行自监督训练 (self-supervised training),PLM 学习蛋白质序列中氨基酸之间的复杂关系和模式。这些模型通常采用Transformer 架构 (Transformer architecture),包含多个Transformer 层 (Transformer layers),每个层都由注意力机制 (attention mechanisms)前馈神经网络 (feed-forward neural networks) 组成,逐步构建蛋白质序列的中间表示 (representations)。PLM 能够捕捉蛋白质的共同进化模式,以及与蛋白质结构和功能相关的物理化学性质。例如,ESM-2 (Evolutionary Scale Modeling-2) 是 Meta AI 开发的一个大型 PLM,它在未标记的蛋白质序列上进行预训练,能够生成高质量的蛋白质嵌入 (embeddings),这些嵌入已被证明对预测蛋白质结构(如 ESMFold 和 AlphaFold)和功能非常有效。

3.1.2. 嵌入 (Embeddings)

自然语言处理 (Natural Language Processing, NLP) 和生物信息学领域中,嵌入 (embeddings) 是指将离散的符号(如单词、氨基酸)映射到连续的、低维向量空间中的表示。这些向量捕捉了符号之间的语义和上下文关系。在 PLM 中,输入的蛋白质序列中的每个氨基酸(或词元 (token))都会被转换成一个嵌入向量 (embedding vector)。随着信息流经 Transformer 模型的不同层,这些嵌入向量会被不断更新和精炼,编码了越来越抽象和高级的蛋白质信息。这些中间层的嵌入是本文中 SAE 训练的输入。

3.1.3. 稀疏自编码器 (Sparse Autoencoders, SAEs)

稀疏自编码器 (Sparse Autoencoders, SAEs) 是一种特殊类型的神经网络 (neural network),旨在学习输入数据的一种稀疏表示。其核心思想是,编码器 (encoder) 将输入数据压缩成一个低维的潜在表示 (latent representation),然后解码器 (decoder) 尝试从这个潜在表示中重建原始输入。关键在于“稀疏”:SAE 通过在潜在层强制稀疏性 (sparsity) 约束(例如,通过添加 L1 正则化项),鼓励模型学习到尽可能少的神经元被激活来表示输入。这使得学习到的潜在特征更具解耦性 (disentangled)可解释性 (interpretability),因为每个潜在特征往往对应于输入数据的特定方面或模式。在本文中,SAE 被用于从 PLM 的神经元激活中解耦出更原子化、更可解释的生物学概念。

3.1.4. 叠加 (Superposition)

叠加 (Superposition) 是指在神经网络中,多个独立的特征或概念被编码在同一组神经元中,而不是每个特征独占一个神经元。这通常发生在神经网络的维度不足以容纳所有需要表示的特征时。例如,一个神经元可能同时对蛋白质的结合位点和某种结构基序都有响应。这种现象使得直接分析单个神经元来理解模型的内部工作机制变得困难,因为单个神经元可能具有多义性 (polysemantic)。SAE 的目标就是通过引入稀疏性来解决叠加问题,将这些叠加的概念解耦 (decouple) 成单独的稀疏特征。

3.1.5. Swiss-Prot 数据库 (Swiss-Prot Database)

Swiss-ProtUniProtKB (Universal Protein Resource Knowledgebase) 的一个子集,是世界上最全面、高质量且手工注释的蛋白质序列和功能信息数据库。它提供了关于蛋白质序列、功能、结构、翻译后修饰、相互作用、遗传疾病等方面的详细信息。在本文中,Swiss-Prot 的注释被用作真实标注数据 (Ground Truth),以量化评估 SAE 提取的特征与已知生物学概念的关联程度,从而衡量这些特征的生物学可解释性 (interpretability)

3.1.6. 大语言模型 (LLM)

大语言模型 (Large Language Models, LLMs) 是指具有数亿甚至数千亿参数的深度学习模型,它们在海量的文本数据上进行训练,能够理解、生成和处理人类语言。LLM 展现出强大的零样本 (zero-shot)少样本 (few-shot) 学习能力,可以执行翻译、摘要、问答等多种自然语言处理 (NLP) 任务。在本文中,LLM (具体是 Claude-3.5 Sonnet) 被创造性地用于自动解释 SAE 发现的蛋白质特征,通过分析特征的激活模式和相关蛋白质元数据,生成可读的自然语言描述。

3.2. 前人工作

3.2.1. PLM 可解释性研究

早期的 PLM 可解释性研究主要集中在分析 Transformer 模型中的注意力权重 (attention weights),以识别序列中重要的相互作用位点或氨基酸残基 [6]。另一些工作则尝试识别模型中与特定功能(如结合口袋)相关的神经元 (neurons) [7]。这些研究揭示了 PLM 确实学习到了与蛋白质物理原理相关的模式,例如共同进化模式、偶联序列模式和结构基序 [8, 9]。然而,这些方法通常受限于神经元的多义性 (polysemanticity)叠加 (superposition) 现象,使得从单个神经元中提取清晰的、原子化的生物学概念变得困难。本文通过引入 SAE 来解决这一核心挑战。

3.2.2. 稀疏自编码器 (SAE) 在可解释性领域的应用

稀疏自编码器 (SAEs) 已经被成功应用于自然语言处理 (NLP)计算机视觉 (Computer Vision) 领域,以解释大型语言模型 (LLMs)视觉模型 (vision models) 的内部机制 [10, 11, 13, 15]。例如,SAE 曾被用于分析 LLM 中单个神经元的多义性 (polysemanticity),揭示了多个不相关的概念可能被编码在同一神经元中 [10, 11]。SAE 的稀疏性 (sparsity) 约束有助于将这些叠加的概念解耦 (disentangle) 成更原子化、更可解释的特征。近期工作还利用 SAE 激活来推断模型中复杂行为的负责任表示 [17],并成功分析了 Transformer 中的电路 (circuits) [16]。本文将 SAE 的成功经验扩展到生物领域,以解释 PLM 的内部机制。

3.3. 技术演进

PLM 领域从最初基于隐马尔可夫模型 (HMMs)进化统计 (evolutionary statistics) 的方法,逐步演进到利用深度学习 (deep learning),特别是Transformer 架构 (Transformer architecture) 的模型。这些模型通过自监督学习 (self-supervised learning) 从大规模的蛋白质序列数据中捕捉共同进化关系和物理化学性质。随着模型规模的不断扩大,其预测能力也显著提升,但可解释性 (interpretability) 的挑战也日益突出。

在可解释性领域,研究从最初对模型输入-输出的归因分析 (attribution analysis),发展到对模型内部组件(如注意力权重、单个神经元)的直接探究。SAE 的引入代表了在解决神经元叠加 (superposition) 问题上的一个重要进展,它提供了一种将复杂、多义的神经元激活分解为更简单、原子化特征的机制。本文将这一先进的可解释性技术引入 PLM 领域,标志着 PLM 可解释性研究向更深层次的机制可解释性 (mechanistic interpretability) 迈进。

3.4. 差异化分析

本文的方法与现有工作的主要区别和创新点在于:

  • 克服叠加问题: 现有 PLM 可解释性研究主要受限于神经元的多义性 (polysemanticity)叠加 (superposition),导致难以提取清晰的生物学概念。本文通过引入 SAE 明确解决了这一问题,成功地将 PLM 中的叠加概念解耦 (decouple) 为稀疏、可解释的特征,从而实现了比直接分析神经元显著更高的概念覆盖度和清晰度。
  • 系统性与定量评估: 本文提出了一套完整的系统性框架,包括 SAE 训练、特征-概念关联分析、基于 LLM 的自动解释以及实际应用验证。并利用 Swiss-Prot 数据库 (Swiss-Prot database) 进行了严格的定量评估 (quantitative evaluation),证明了 SAE 特征在捕捉已知生物学概念方面的优越性。这为 PLM 可解释性研究设定了新的标准。
  • 发现新颖生物学: 不同于仅仅验证模型对已知知识的编码,本文还展示了 PLM 学习到了超出现有注释连贯概念 (coherent concepts)。这为利用 PLM 进行新生物学发现开辟了途径。
  • 自动化解释与应用: 创新性地将大语言模型 (LLM) 引入特征解释流水线,实现了对大量特征的自动解释 (automatic interpretation),极大地提升了可扩展性。此外,本文不仅停留在解释层面,还展示了这些可解释特征在填充缺失注释 (filling missing annotations)定向引导蛋白质生成 (targeted steering of protein generation) 等实际应用中的潜力。

4. 方法论

本文的核心方法是利用稀疏自编码器 (Sparse Autoencoders, SAEs)蛋白质语言模型 (PLM) 的中间层嵌入 (embeddings) 中提取稀疏 (sparse)可解释 (interpretable) 的特征。随后,这些特征通过与已知生物学概念进行关联分析、利用大语言模型 (LLM) 自动生成描述以及在实际应用中进行验证。

4.1. 方法原理

该方法的核心思想 (core idea) 是,PLM 内部的神经元可能以叠加 (superposition) 的方式编码多个生物学概念,导致单个神经元难以直接解释。SAE 通过在潜在层 (latent layer) 强制稀疏性 (sparsity) 约束,能够将这些叠加的概念解耦 (disentangle) 为更原子化、更易于理解的稀疏特征。每个稀疏特征理论上应该对应一个更具体的生物学概念或模式。

具体来说,SAE 学习一个字典 (dictionary),其中每个字典向量 (dictionary vector) 代表一个潜在特征。当 PLM 的一个中间层嵌入输入到 SAE 时,SAE 会激活少数几个字典向量来重构该嵌入。这些被激活的字典向量及其激活值就构成了该输入的稀疏特征表示。通过分析这些稀疏特征的激活模式,并将其与蛋白质的生物学注释相关联,可以揭示 PLM 内部学习到的生物学原理。

下图(原文 Figure 1)展示了 SAE 方法的概述以及通过自动激活分析揭示的代表性 SAE 特征:

该图像是示意图,展示了使用稀疏自编码器(SAE)从蛋白质语言模型中提取可解释特征的步骤以及激活与已知生物概念相关的特征。图中包括与结构接近性和顺序接近性相关的特征激活情况,以及与核心结合域和二硫键相关的激活示例。 图1:SAE 方法的概述以及通过自动激活分析揭示的代表性 SAE 特征。左侧面板(a)展示了从蛋白质语言模型中提取稀疏自编码器特征的流程。右侧面板展示了通过激活模式与已知 Swiss-Prot 生物概念(如核心结合域和二硫键)显著关联而选择的代表性 SAE 特征。

4.2. 核心方法详解 (逐层深入)

4.2.1. 稀疏自编码器训练 (Sparse Autoencoder Training)

4.2.1.1. 数据集准备 (Dataset Preparation)

为了训练 SAE,首先需要准备用于提取 PLM 隐藏表示 (hidden representations) 的蛋白质序列数据集。

  • 选择 PLM: 作者选择了 ESM-2-8M-UR50D 模型。这是一个具有 8 亿参数的 ESM-2 模型变体,在 UniRef50 数据集上预训练。
  • 数据集来源:UniRef50 数据集的一个子集随机选择了 1000 万个蛋白质序列。
  • 嵌入提取: 从 ESM-2-8M 模型的Transformer 块 (transformer block) 的第 1 到第 6 层中提取隐藏表示 (hidden representations)。在提取过程中,排除了特殊的 <cls><cls> (分类 token) 和 <eos><eos> (序列结束 token) 的词元 (tokens),因为 SAE 主要关注序列中氨基酸的特征。
  • 数据打乱: 在每个训练阶段之前,对数据集进行打乱,以确保训练过程中数据的随机采样。

4.2.1.2. 架构和训练参数 (Architecture and Training Parameters)

SAE 的架构包括一个编码器 (encoder) 和一个解码器 (decoder)

  • 维度扩展: ESM-2 模型的嵌入向量 (embedding vectors) 大小为 320 维。SAE 将每层的神经元从 320 个扩展到 10,240 个潜在特征 (latent features),这意味着学习到的字典大小 ddictd_{\mathrm{dict}} 为 10,240。

  • 训练过程: 对每个 ESM-2 层,训练了 20 个 SAE,每个 SAE 训练 500,000 步。

  • 批次大小 (Batch Size): 2,048。

  • 学习率 (Learning Rate): 从 1e-4 到 1e-8,以 10 倍的增量采样。

  • L1 正则化 (L1 Regularization):

    • L1 正则化项 λ\lambda 的值在 0.07 到 0.1 之间。L1 正则化是强制稀疏性 (sparsity) 的关键。它通过向损失函数添加潜在特征 (latent features) 激活值的绝对值之和,惩罚大的激活值,从而鼓励模型学习只激活少量特征来表示输入。
  • 学习率调度 (Learning Rate Schedule): 在训练的最初 5% 阶段,学习率达到其最大值,然后逐渐衰减。

    SAE 的重建 (reconstruction) 过程可以表示为: xb+i=1ddictfi(x)di \mathbf { x } \approx \mathbf { b } + \sum _ { i = 1 } ^ { d _ { \mathrm { d i c t } } } f _ { i } ( \mathbf { x } ) \mathbf { d } _ { i } 其中:

  • x\mathbf { x } 是输入的激活向量(来自 ESM-2 的嵌入)。

  • b\mathbf { b } 是一个偏置项。

  • ddictd _ { \mathrm { d i c t } } 是学习到的字典大小(即潜在特征的数量)。

  • fi(x)f _ { i } ( \mathbf { x } ) 是第 ii潜在特征 (latent feature) 的激活值,表示该特征对重构的贡献。

  • di\mathbf { d } _ { i } 是第 ii 个潜在特征对应的字典向量 (dictionary vector),存储在解码器矩阵 (decoder matrix) 的行中,它定义了该特征在原始嵌入空间中的“形状”。

    SAE 的编码器和解码器通常定义如下: 编码器计算特征激活 fi(x)f_i(x)fi(x)=ReLU(Wex+be)i f_i(\mathbf{x}) = \mathrm{ReLU}(W_e \mathbf{x} + b_e)_i 其中:

  • WeW_e编码器权重矩阵 (encoder weight matrix)

  • beb_e编码器偏置向量 (encoder bias vector)

  • ReLU\mathrm{ReLU}修正线性单元 (Rectified Linear Unit) 激活函数,确保激活值为非负。

    解码器重构输入 x\mathbf{x}': x=Wdf(x)+bd \mathbf{x}' = W_d f(\mathbf{x}) + b_d 其中:

  • WdW_d解码器权重矩阵 (decoder weight matrix),其行即为字典向量 di\mathbf{d}_i

  • bdb_d解码器偏置向量 (decoder bias vector)

    下图(原文 Figure 8)展示了 SAE 的分解和训练概览:

    Figure 8: Overview of SAE decomposition and training. (a) Decomposition of embedding vector into weighted sum of dictionary elements.) Architecture for the SAE 图8:SAE 的分解和训练概览。(a) 将嵌入向量分解为字典元素的加权和。(b) SAE 的架构。

4.2.1.3. 特征归一化 (Feature Normalization)

为了标准化特征比较,作者使用 Z-score 归一化 (Z-score normalization) 对激活值进行归一化。这通过对来自 5,000 个蛋白质的激活值进行扫描 (scan) 来实现,确保所有特征的激活值具有可比的尺度。

4.2.2. Swiss-Prot 概念评估流水线 (Swiss-Prot Concept Evaluation Pipeline)

为了量化评估 SAE 提取特征的生物学可解释性 (interpretability),作者建立了一个将特征激活与已知 Swiss-Prot 概念 (Swiss-Prot concepts) 关联起来的流水线。

4.2.2.1. 数据集构建 (Dataset Construction)

  • UniProtKB (Swiss-Prot) 的审查子集中,随机采样了 50,000 个长度小于 1,024 个氨基酸的蛋白质。
  • 将这些蛋白质分为训练集、验证集和测试集,比例分别为 80%、10% 和 10%。
  • 验证集和测试集中的蛋白质至少包含 1,500 个氨基酸,以确保有足够的数据进行评估。

4.2.2.2. 特征-概念关联分析 (Feature-Concept Association Analysis)

  • 激活阈值: 对于每个 SAE 特征,通过尝试一系列激活阈值 (activation thresholds)(0.1、0.2、0.4、0.6、0.8)来确定其与 Swiss-Prot 概念的关联。
  • 评估指标: 使用修改后的精确率 (precision)召回率 (recall) 指标来评估特征-概念关联。 precision=TruePositivesTruePositives+FalsePositivesrecall=DomainsWithTruePositiveTotalDomainsF1=2precisionrecallprecision+recall \begin{array} { r } { \mathrm { precision } = \frac { \mathrm { TruePositives } } { \mathrm { TruePositives } + \mathrm { FalsePositives } } } \\ { \mathrm { recall } = \frac { \mathrm { DomainsWithTruePositive } } { \mathrm { TotalDomains } } } \\ { \mathrm { F1 } = 2 \cdot \frac { \mathrm { precision } \cdot \mathrm { recall } } { \mathrm { precision } + \mathrm { recall } } } \end{array} 其中:
    • TruePositives\mathrm{TruePositives} (真阳性):特征激活且蛋白质具有该概念注释的氨基酸位置数量。
    • FalsePositives\mathrm{FalsePositives} (假阳性):特征激活但蛋白质不具有该概念注释的氨基酸位置数量。
    • DomainsWithTruePositive\mathrm{DomainsWithTruePositive} (具有真阳性的域):至少在一个位置上特征激活且蛋白质具有该概念注释的蛋白质域数量。
    • TotalDomains\mathrm{TotalDomains} (总域数):数据集中具有该概念注释的蛋白质域总数量。
    • F1\mathrm{F1} (F1 分数):精确率和召回率的调和平均值 (harmonic mean),综合衡量模型的性能。
  • 最佳 F1 分数: 对于每个特征-概念对,选择在验证集上产生最高 F1 分数的激活阈值进行最终评估。

4.2.2.3. 模型选择和评估 (Model Selection and Evaluation)

  • 概念筛选: 仅对在验证集中拥有超过 10 个域或 1,500 个氨基酸的 135 个概念进行初步评估。
  • 模型选择: 对于每个 ESM-2-8M 层,选择 F1 分数最高的 SAE 作为后续分析和 InterPLM 仪表板中使用的模型。
  • 最终评估: 识别在验证集中 F1 分数大于 0.5 的特征-概念对,然后在测试集上计算它们的 F1 分数,并报告在测试集中 F1 分数仍大于 0.5 的数量。

4.2.2.4. 基线 (Baselines)

为了比较,作者训练了两种基线模型:

  • ESM-2 神经元: 直接使用 ESM-2 模型的原始神经元激活作为特征,并按照与 SAE 特征相同的方式进行评估。
  • 随机 ESM-2 权重: 训练一个具有随机初始化权重的 ESM-2 模型,然后在此模型上训练 SAE。这用于评估 SAE 捕捉生物学概念的能力是否依赖于 PLM 预训练学习到的有意义表示,而不是模型架构本身。

4.2.3. LLM 特征注释流水线 (LLM Feature Annotation Pipeline)

为了超越 Swiss-Prot 概念的限制,作者开发了一个利用大语言模型 (LLM) 自动生成特征描述的流水线。

4.2.3.1. 示例选择 (Example Selection)

  • 对 1,200 个 (10%) 随机选择的特征进行分析。
  • 对于每个特征,通过选择激活值差异最大的蛋白质(即具有高、中、低激活值的蛋白质)来选择代表性蛋白质 (representative proteins)
  • 激活值被量化到 10 个区间(例如,0-0.1,0.1-0.2,...,0.9-1.0)。
  • 从每个区间选择蛋白质,以提供多样化的激活示例。

4.2.3.2. 描述生成和验证 (Description Generation and Validation)

  • 输入 LLM:Claude-3.5 Sonnet (新的) 提供特征 (feature) 的 Swiss-Prot 元数据 (metadata)定量激活值 (quantitative activation values),包括激活氨基酸及其在蛋白质中的位置。

  • 生成描述: LLM 被提示生成对特征激活模式的描述,以及在给定描述和蛋白质元数据的情况下,能够预测新蛋白质的激活值的摘要。描述的开头必须是:“激活模式的特点是: (The activation patterns are characterized by:)”。摘要的开头必须是:“该特征激活于: (The feature activates on...)”。

  • 描述验证: 为了验证 LLM 生成描述的准确性,向 Claude 提供了依赖于描述的蛋白质最大激活值 (maximum activation value) 的预测。然后,将这些预测值与实际测量值 (measured values) 进行比较,使用皮尔逊相关系数 (Pearson correlation) 进行量化,以评估描述的预测能力。

    下图(原文 Figure 5)展示了 LLM 生成自动特征描述的工作流程:

    该图像是示意图,展示了如何生成和评估蛋白质语言模型特征描述,左侧部分描述了利用SwissProt元数据生成特征描述的流程,右侧展示了特征激活值的Pearson相关系数的分布,底部展示了特征8386、10091和7404的具体示例及其对应的预测激活与真实激活的关系。 图5:语言模型可以为 SAE 特征生成自动特征描述。(a) 使用 Claude-3.5 Sonnet (new) 生成和验证描述的工作流程。(b) 比较生成的特征描述的最大激活值与蛋白质预测值(通过核密度估计可视化)。文本是 Claude 对每个特征元素描述的总结,并在结构旁边标注了最大激活示例中的描述。

4.2.4. 特征分析和可视化 (Feature Analysis and Visualization)

  • UMAP 嵌入和聚类 (UMAP Embedding and Clustering):

    • 使用 UMAP (Uniform Manifold Approximation and Projection) 对归一化的 SAE 解码器权重 (decoder weights) 进行降维 (dimensionality reduction),以便在二维空间中可视化特征。UMAP 参数为:metric=euclideanmetric='euclidean', neighbors=15neighbors=15, min_dist=0.1
    • 使用 HDBSCAN (Hierarchical Density-Based Spatial Clustering of Applications with Noise) 算法对 UMAP 嵌入的特征进行聚类 (clustering),以便在 InterPLM 界面中进行可视化。HDBSCAN 参数为:min_cluster_size=5, min_samples=3
  • 序列和结构特征分析 (Sequential and Structural Feature Analysis):

    • 识别高激活区域: 识别蛋白质中激活值高于 0.6 的区域,这些蛋白质具有可用的 AlphaFold 结构 (AlphaFold structures)
    • 计算序列和结构聚类: 对于每个蛋白质的最高激活残基,计算:
      • 序列聚类 (Sequential clustering): 在序列中 ±2\pm 2 个位置内的平均激活值。
      • 结构聚类 (Structural clustering): 在 3D 空间中 6 Å 范围内的残基平均激活值。
    • 生成零分布 (Null Distributions): 通过对每个蛋白质进行 5 次随机排列并取平均值来生成零分布,以作为比较的基线。
    • 评估聚类显著性: 使用配对 t 检验 (paired t-tests)Cohen's d 效应大小 (Cohen's d effect sizes) 对 100 个蛋白质评估聚类显著性。筛选出 Bonferroni 校正后的结构 p 值小于 0.05 的蛋白质,并根据结构与序列效应大小的比值对特征进行着色。

4.2.5. 引导实验 (Steering Experiments)

为了演示可解释特征的实际应用,作者进行了序列引导 (sequence steering) 实验。

  • 分解嵌入: 遵循 [20] 中描述的方法,将 ESM 嵌入 (embeddings) 分解为 SAE 重建预测 (reconstruction predictions)误差项 (error terms)

  • 引导流程:

    1. 提取指定层的嵌入。
    2. 计算 SAE 重建和误差。
    3. 通过激活或去激活所需的特征来修改重建。
    4. 将修改后的重建与误差项结合。
    5. 允许模型进行正常的处理。
    6. 使用 NNsight 工具提取预测的嵌入。
  • 目标: 展示激活特定特征如何引导 (steer) 甘氨酸 (Glycine) 在周期性模式(例如胶原蛋白样区域中的 GxX 重复)中的预测。

    下图(原文 Figure 10)展示了在引导周期性甘氨酸特征时 P(Glycine) 的变化:

    该图像是示意图,展示了在不同的“steer”参数下,蛋白质序列中各个位置的概率(G)。图中上部分为柱状图,表示不同“steer”值下的概率分布;下部分展示蛋白质的结构形式,支持对序列生成过程的可视化理解。不同颜色标识不同的“steer”强度。整体上,该图表明了调整“steer”参数对序列生成的影响。 图10:在引导周期性甘氨酸特征时 P(Glycine) 的变化。

下图(原文 Figure 11)展示了在引导非周期性甘氨酸特征时 P(Glycine) 的变化:

该图像是一个条形图和结构示意图,展示了在不同引导力度下,某蛋白质序列位置的概率分布。条形图中,各颜色代表不同的引导力度,从无引导到2.5倍引导,显示了在特定序列位置上各状态的概率变化。图下方为该序列的空间位置结构示意。 图11:在引导非周期性甘氨酸特征时 P(Glycine) 的变化。

5. 实验设置

5.1. 数据集

5.1.1. SAE 训练数据集

  • 来源:UniRef50 数据集的一个子集随机选择了 1000 万个蛋白质序列。UniRef50 是一个非冗余蛋白质序列数据库,其中序列之间的相似性限制在 50% 以下,有助于确保训练数据的多样性。
  • 特点: 这些蛋白质用于从 ESM-2-8M-UR50D 模型中提取隐藏表示 (hidden representations),作为 SAE 的输入。

5.1.2. Swiss-Prot 概念评估数据集

  • 来源:UniProtKB (Swiss-Prot) 的审查子集中,随机采样了 50,000 个蛋白质。Swiss-Prot 是一个高质量、手工注释的蛋白质序列和功能信息数据库,提供了丰富的生物学概念注释。
  • 规模与特点:
    • 所有蛋白质长度均小于 1,024 个氨基酸。
    • 验证集和测试集中的蛋白质至少包含 1,500 个氨基酸。
    • 数据集被划分为 80% 训练集、10% 验证集和 10% 测试集。
  • 用途: 用于量化 SAE 提取的特征与已知 Swiss-Prot 概念的关联程度,从而评估其生物学可解释性 (interpretability)

5.1.3. AlphaFold 结构数据

  • 用途:特征分析和可视化 (Feature Analysis and Visualization) 部分,用于识别蛋白质的高激活区域 (high-activation regions),并通过 AlphaFold 预测的 3D 结构来计算结构聚类 (structural clustering) 特性。

5.2. 评估指标

本文使用了多种评估指标来衡量 SAE 特征的性能和可解释性 (interpretability)

5.2.1. 精确率 (Precision)

  1. 概念定义 (Conceptual Definition): 精确率衡量的是在所有被模型预测为正类(即特征激活)的样本中,有多少比例确实是正类(即具有生物学概念注释)。它关注的是模型预测为正类的结果的准确性。
  2. 数学公式 (Mathematical Formula): precision=TruePositivesTruePositives+FalsePositives \mathrm { precision } = \frac { \mathrm { TruePositives } } { \mathrm { TruePositives } + \mathrm { FalsePositives } }
  3. 符号解释 (Symbol Explanation):
    • TruePositives\mathrm{TruePositives}真阳性 (True Positives),指特征激活且蛋白质在该位置具有目标生物学概念注释的氨基酸数量。
    • FalsePositives\mathrm{FalsePositives}假阳性 (False Positives),指特征激活但蛋白质在该位置不具有目标生物学概念注释的氨基酸数量。

5.2.2. 召回率 (Recall)

  1. 概念定义 (Conceptual Definition): 召回率衡量的是在所有实际为正类(即具有生物学概念注释)的样本中,有多少比例被模型正确地预测为正类。它关注的是模型识别出所有相关正类的能力。
  2. 数学公式 (Mathematical Formula): recall=DomainsWithTruePositiveTotalDomains \mathrm { recall } = \frac { \mathrm { DomainsWithTruePositive } } { \mathrm { TotalDomains } }
  3. 符号解释 (Symbol Explanation):
    • DomainsWithTruePositive\mathrm{DomainsWithTruePositive}:指在给定数据集中,至少在一个位置上,特征激活且蛋白质具有目标生物学概念注释的蛋白质域数量。
    • TotalDomains\mathrm{TotalDomains}:指在给定数据集中,所有具有目标生物学概念注释的蛋白质域总数量。

5.2.3. F1 分数 (F1 Score)

  1. 概念定义 (Conceptual Definition): F1 分数是精确率和召回率的调和平均值 (harmonic mean)。它综合考虑了精确率和召回率,因此在分类任务中,尤其是在类别不平衡的情况下,是一个更稳健的评估指标。F1 分数越高,表示模型在精确性和完整性之间达到了更好的平衡。
  2. 数学公式 (Mathematical Formula): F1=2precisionrecallprecision+recall \mathrm { F1 } = 2 \cdot \frac { \mathrm { precision } \cdot \mathrm { recall } } { \mathrm { precision } + \mathrm { recall } }
  3. 符号解释 (Symbol Explanation):
    • precision\mathrm{precision}精确率 (precision),如上所述。
    • recall\mathrm{recall}召回率 (recall),如上所述。

5.2.4. L0 范数 (L0 Norm)

  1. 概念定义 (Conceptual Definition): L0 范数衡量的是一个向量中非零元素的数量。在稀疏自编码器 (SAE) 的上下文中,它用于量化潜在表示 (latent representation) 的稀疏性,即对于一个给定的输入,有多少个潜在特征被激活(其激活值大于零)。L0 范数越小,表示稀疏性越高。
  2. 数学公式 (Mathematical Formula): L0(f(x))=i=1ddetanor1(fi(x)>0) L _ { 0 } ( f ( x ) ) = \displaystyle \sum _ { i = 1 } ^ { d _ { \mathrm { d e t a n o r } } } \mathbf { 1 } ( f _ { i } ( x ) > 0 )
  3. 符号解释 (Symbol Explanation):
    • f(x):输入的潜在表示向量。
    • fi(x)f_i(x):向量 f(x) 的第 ii 个元素(即第 ii 个潜在特征的激活值)。
    • ddetanord_{\mathrm{detanor}}:潜在特征的总数量(即字典大小)。
    • 1()\mathbf{1}(\cdot)指示函数 (indicator function),如果括号内的条件为真,则返回 1,否则返回 0。

5.2.5. 均方误差 (Mean Squared Error, MSE)

  1. 概念定义 (Conceptual Definition): 均方误差衡量的是模型预测值与真实值之间差异的平方的平均值。在稀疏自编码器 (SAE) 的上下文中,它用于评估 SAE 重建其输入嵌入 (embeddings) 的准确性。MSE 值越低,表示重建质量越高。
  2. 数学公式 (Mathematical Formula): MSE(x,x)=1ddetionori=1dmodet(xixi)2 M S E ( x , x ^ { \prime } ) = \displaystyle \frac { 1 } { d _ { \mathrm { d e t i o n o r } } } \sum _ { i = 1 } ^ { d _ { \mathrm { m o d e t } } } ( x _ { i } - x _ { i } ^ { \prime } ) ^ { 2 }
  3. 符号解释 (Symbol Explanation):
    • xx:原始输入向量(即 ESM-2 模型的嵌入)。
    • xx':SAE 重建的输出向量。
    • xix_i:原始输入向量的第 ii 个元素。
    • xix_i':重建输出向量的第 ii 个元素。
    • ddetionord_{\mathrm{detionor}}:输入向量的维度。

5.2.6. L1 范数 (L1 Norm)

  1. 概念定义 (Conceptual Definition): L1 范数衡量的是一个向量中所有元素的绝对值之和。在稀疏自编码器 (SAE) 的训练中,L1 范数通常作为正则化项 (regularization term) 添加到损失函数中,以强制稀疏性 (sparsity)。惩罚大的激活值,从而鼓励模型学习只激活少量特征来表示输入。
  2. 数学公式 (Mathematical Formula): L1(f(x))=i=1ddetanorfi(x) L _ { 1 } ( f ( x ) ) = \displaystyle \sum _ { i = 1 } ^ { d _ { \mathrm { d e t a n o r } } } | f _ { i } ( x ) |
  3. 符号解释 (Symbol Explanation):
    • f(x):输入的潜在表示向量。
    • fi(x)f_i(x):向量 f(x) 的第 ii 个元素(即第 ii 个潜在特征的激活值)。
    • ddetanord_{\mathrm{detanor}}:潜在特征的总数量(即字典大小)。
    • |\cdot|:绝对值函数。

5.2.7. 损失恢复百分比 (% Loss Recovered)

  1. 概念定义 (Conceptual Definition): 损失恢复百分比衡量的是 SAE 重建的嵌入 (embeddings) 在多大程度上能够恢复原始模型(ESM-2)的交叉熵损失 (cross-entropy loss)。它是评估 SAE 重建质量的一个重要指标,反映了 SAE 捕捉原始嵌入中关键信息的能力。高恢复百分比意味着 SAE 能够有效地捕获原始嵌入中对模型预测有用的信息。
  2. 数学公式 (Mathematical Formula): GL Loss Recovered=(1CEReconstumCEOniginalCEZeeonCEOniginal)×100 \displaystyle \mathcal { G } _ { L } \mathrm { ~ L o s s ~ R e c o v e r e d } = \Big ( 1 - \frac { C E _ { \mathrm { Reconstum } } - C E _ { \mathrm { Oniginal } } } { C E _ { \mathrm { Zeeon } } - C E _ { \mathrm { Oniginal } } } \Big ) \times 100
  3. 符号解释 (Symbol Explanation):
    • CEReconstumCE_{\mathrm{Reconstum}}:使用 SAE 重建的嵌入时模型的交叉熵损失。
    • CEOriginalCE_{\mathrm{Original}}:使用原始 ESM-2 嵌入时模型的交叉熵损失。
    • CEZeeonCE_{\mathrm{Zeeon}}:当指定嵌入层被替换为全零(但后续层保持不变)时模型的交叉熵损失。

5.3. 对比基线

本文将 SAE 提取的特征与以下基线进行了比较:

  • ESM-2 神经元 (ESM-2 Neurons): 这是最直接的基线,通过直接分析 ESM-2 模型的原始神经元激活来评估其可解释性 (interpretability)。这有助于量化 SAE 在解决神经元叠加 (superposition) 问题上的优势。
  • 随机 ESM-2 权重训练的 SAE (SAE on Randomized ESM-2 Weights): 通过在具有随机初始化权重的 ESM-2 模型上训练 SAE,来评估 SAE 捕捉生物学概念的能力是否依赖于 PLM 预训练学习到的有意义表示。如果这种随机模型的 SAE 也能捕获生物学概念,则可能表明这些概念是数据分布固有的,而非模型通过学习获得的。

6. 实验结果与分析

6.1. 核心结果分析

6.1.1. 稀疏自编码器 (SAE) 在蛋白质语言模型 (PLM) 中发现可解释概念

本文的核心发现是,SAE 能够从 ESM-2 的嵌入中提取出大量人类可解释 (human-interpretable)潜在特征 (latent features)。通过将 ESM-2 的六个层从 320 个神经元扩展到 10,420 个潜在特征,SAE 成功地揭示了与生物学概念高度关联的激活模式(例如,催化位点、结合位点、蛋白质域)。

通过量化评估,SAE 特征比 ESM-2 神经元捕捉到显著更多概念关联 (conceptual alignment)。例如,SAE 特征与 Swiss-Prot 概念的 F1 分数显著高于 ESM 神经元。即使在仅有少量已知注释的生物学概念上,SAE 特征也表现出卓越的性能。

下图(原文 Figure 3)展示了 SAE 特征与 ESM 神经元在 Swiss-Prot 概念关联方面的比较:

该图像是一个示意图,展示了不同层次的F1分数在多个类别下的分布,层数从1到6,F1分数范围在0到1之间,不同颜色的点代表不同的类别。 图3:SAE 特征与 ESM 神经元相比,与 Swiss-Prot 概念具有更强的关联。(a) 每个 Swiss-Prot 概念的最高 F1 分数。比较了在 ESM-2 嵌入上训练的 SAE 特征(粉色)、ESM-2 嵌入中的神经元(蓝色)以及在具有随机权重的 ESM-2 模型嵌入上训练的 SAE 特征(绿色)的性能。模型根据特征与 Swiss-Prot 概念之间的 F1 分数进行比较。(b) 对于每个层,计算在验证集和测试集中 F1 分数均大于 0.5 的特征数量。

从 Figure 3 可以看出,粉色条(SAE 特征)在大多数概念上都达到了更高的 F1 分数,并且具有高 F1 分数的特征数量远多于蓝色条(ESM 神经元)。这表明 SAE 有效地解耦 (disentangled) 了 PLM 内部的叠加 (superposition) 概念,使其变得更加可解释。随机 ESM 的 SAE(绿色条)在生物学概念上几乎没有关联,这进一步证实了 SAE 捕捉的是 PLM 预训练学习到的有意义的生物学信息,而非模型架构本身。

下图(原文 Figure 5)展示了每层中人类可解释特征的数量:

该图像是一个柱状图,展示不同层级(1到6层)中人类可解释特征的数量。粉色柱子表示通过稀疏自编码器(SAE)提取的特征数量,蓝色柱子则标识了神经元的数量,绿色柱子显示随机ESM的特征,此图揭示了特征提取的效率随着层数的增加而变化。 图5:SAE 特征与 ESM 神经元相比,与 Swiss-Prot 概念具有更强的关联。粉色柱子表示通过稀疏自编码器(SAE)提取的特征数量,蓝色柱子则标识了神经元的数量,绿色柱子显示随机ESM的特征,此图揭示了特征提取的效率随着层数的增加而变化。

结果显示,SAE 特征每层能够揭示多达 2,548 个清晰的概念关联,而 ESM 神经元每层最多只有 46 个。这证实了 PLM 中存在叠加 (superposition) 现象,且 SAE 能够有效解决这一问题。

6.1.2. 特征基于共享功能和结构角色形成聚类

SAE 特征不仅在个体层面具有可解释性 (interpretability),它们还根据共享的生物学功能和结构角色形成有意义的聚类 (clusters)。例如,作者观察到与激酶 (kinase) 活性相关的特征,它们集中在催化环附近的不同区域,包括催化环前的 β 折叠。尽管这些特征的峰值激活位置 (peak activation positions) 不同,但它们在各自的最大激活示例上都保持高激活水平,表明它们识别相似的激酶亚型。

另一个例子是TonB 依赖性受体 (TBDR) 的 β 桶结构。某些特征表现出高度特异性(如 f/1503,F1=0.998),作为真正的 TBDR 检测器,而另一些则是更广义的 β 桶特征。这表明 SAE 能够学习到不同粒度(从高度特异性到广义)的结构和功能模式。

下图(原文 Figure 12)展示了 SAE 特征分析和可视化:

Figure 8: Overview of SAE decomposition and training. (a) Decomposition of embedding vector into weighted sum of dictionary elements.) Architecture for the SAE 图12:SAE 特征分析和可视化揭示了具有多样且一致激活模式的特征。左上(a)是特征激活频率分布;右上(b)是结构与序列激活的平均激活比率,用点大小表示激活频率;左下(c)是特征字典值的 UMAP 聚类;右下(d)是带有 Swiss-Prot 概念标签的特征字典值 UMAP 聚类,突出显示了在糖基转移酶 amsK 蛋白上最大激活的特征。

从 Figure 12c 和 12d 可以看出,SAE 特征在 UMAP 空间中形成清晰的聚类,这些聚类通常与特定的 Swiss-Prot 概念(如糖基转移酶)相关联,进一步支持了 SAE 特征的生物学意义和可解释性 (interpretability)

6.1.3. 大语言模型 (LLM) 可生成有意义的特征描述

由于 Swiss-Prot 概念只能解释不到 20% 的特征,作者开发了一个自动化流水线,利用 Claude-3.5 Sonnet 等大语言模型 (LLMs) 为特征生成描述。通过向 LLM 提供激活模式和相关蛋白质元数据 (metadata),LLM 能够生成高度准确 (accurate) 的特征描述。

下图(原文 Figure 8)展示了 LLM 描述生成与验证的流程和结果:

该图像是示意图,展示了如何生成和评估蛋白质语言模型特征描述,左侧部分描述了利用SwissProt元数据生成特征描述的流程,右侧展示了特征激活值的Pearson相关系数的分布,底部展示了特征8386、10091和7404的具体示例及其对应的预测激活与真实激活的关系。 图8:语言模型可以为 SAE 特征生成自动特征描述。(a) 使用 Claude-3.5 Sonnet (new) 生成和验证描述的工作流程。(b) 比较生成的特征描述的最大激活值与蛋白质预测值(通过核密度估计可视化)。文本是 Claude 对每个特征元素描述的总结,并在结构旁边标注了最大激活示例中的描述。

从 Figure 8b 可以看出,LLM 生成的描述与特征的实际激活值具有较高的皮尔逊相关系数 (Pearson r correlation)(中位数 r = 0.72),这表明 LLM 能够准确捕捉特征的激活模式。即使对于没有 Swiss-Prot 注释的概念(如 f/8386 描述的六肽 β 螺旋),LLM 也能生成连贯的描述,这展示了其在发现新颖生物学 (novel biology) 方面的潜力。

6.1.4. 特征激活识别缺失和新蛋白质注释

SAE 特征的激活模式可以被用来识别蛋白质数据库中缺失的注释 (missing annotations),甚至提出新的蛋白质基序 (new protein motifs)

下图(原文 Figure 9)展示了特征激活模式如何识别缺失和新蛋白质注释:

Figur :Feature activation patterns can be used to identiy missing and new protein annotations.)/939 identies missigmot otation orNudix box. I activates n slemino aciin cnserve positio whi is lbe a sr… 图9:特征激活模式可用于识别缺失和新的蛋白质注释。(a) f/939 识别 Nudix 盒基序的缺失注释。它在保守位置的单个氨基酸上激活,该位置被标记为结构上的金属离子结合位点。右侧是蛋白质示例,其激活以粉色突出显示,并带有 Nudix 标签。左侧蛋白质 (B2GFH1) 在 Swiss-Prot 中没有 Nudix 基序注释,但显示出明显的激活峰(红色突出显示)。这意味着该蛋白质中存在 Nudix 基序。(b) f/3147 识别肽酶的缺失域注释。它在一系列氨基酸模式上激活,这些模式具有粉色突出显示的肽酶激活。左侧蛋白质在 Swiss-Prot 中没有肽酶域注释,但显示出激活峰。这意味着该蛋白质中存在肽酶域。(c) f/9046 建议在细菌糖基转移酶中存在 UDP-N-乙酰-α-D-葡萄糖胺 (UDP-GlcNAc) 和 Mg2+ 结合位点的缺失注释。在这两种结构中,较高的激活值用深粉色表示。右侧蛋白质具有糖基转移酶活性,并有绿色标记的结合位点注释。左侧蛋白质也具有糖基转移酶活性,但缺失这些结合位点注释,其激活峰用粉色突出显示。

  • Nudix 盒基序: 特征 f/939 能够在没有 Nudix 盒注释的蛋白质中被激活(如 B2GFH1),但其激活位置与已知的 Nudix 盒的金属离子结合残基一致,暗示了该蛋白质中可能存在未被标记的 Nudix 盒基序。

  • 肽酶域: 特征 f/3147 在某些未标注肽酶域的蛋白质中显示出高激活,表明这些蛋白质可能包含未被识别的肽酶域。

  • UDP-GlcNAc 和 Mg2+ 结合位点: 特征 f/9046 能够识别细菌糖基转移酶中潜在的 UDP 和 Mg2+ 结合位点,这些位点可能在现有数据库中缺失注释。

    这些结果证明了 SAE 特征作为一种计算工具 (computational tool) 的潜力,可以帮助生物学家补充现有数据库,并加速新生物学发现。

6.1.5. 蛋白质序列生成可通过激活可解释特征进行引导

作者还展示了这些可解释特征 (interpretable features) 可以用于定向引导 (targeted steering) 蛋白质序列生成。通过激活特定特征,可以影响 PLM 在生成序列时对特定氨基酸模式的预测。

下图(原文 Figure 10)展示了在引导周期性甘氨酸特征时 P(Glycine) 的变化:

该图像是示意图,展示了在不同的“steer”参数下,蛋白质序列中各个位置的概率(G)。图中上部分为柱状图,表示不同“steer”值下的概率分布;下部分展示蛋白质的结构形式,支持对序列生成过程的可视化理解。不同颜色标识不同的“steer”强度。整体上,该图表明了调整“steer”参数对序列生成的影响。 图10:在引导周期性甘氨酸特征时 P(Glycine) 的变化。

下图(原文 Figure 11)展示了在引导非周期性甘氨酸特征时 P(Glycine) 的变化:

该图像是一个条形图和结构示意图,展示了在不同引导力度下,某蛋白质序列位置的概率分布。条形图中,各颜色代表不同的引导力度,从无引导到2.5倍引导,显示了在特定序列位置上各状态的概率变化。图下方为该序列的空间位置结构示意。 图11:在引导非周期性甘氨酸特征时 P(Glycine) 的变化。

实验中,他们测试了三个在胶原蛋白样区域中周期性甘氨酸重复 (periodic glycine repeats)(GxGXX)上激活的特征。通过仅在其中一个 G 位置进行引导并测量对所有后续周期性 G 位置上甘氨酸 (Glycine) 概率的影响,他们发现引导能够显著增强后续周期性 G 位置上甘氨酸的预测似然 (prediction likelihood),并且这种效应可以传播到后续的重复。这表明 SAE 特征能够捕获高阶的序列模式,并且可以用于细粒度 (fine-grained) 地控制蛋白质的序列生成,这在蛋白质设计 (protein design) 任务中具有重要意义。

6.2. 数据呈现 (表格)

以下是原文 Table 1 的结果,展示了各层 SAE 的学习参数和性能指标:

Layer Learning Rate L1 % Loss Recovered
L1 1.0e-7 0.1 99.60921
L2 1.0e-7 0.08 99.31827
L3 1.0e-7 0.1 99.02078
L4 1.0e-7 0.1 98.39785
L5 1.0e-7 0.1 99.32478
L6 1.0e-7 0.09 100

表1:各层学习参数和 SAE 性能指标。

从 Table 1 可以看出,SAE 在所有层上都实现了非常高的损失恢复百分比 (% Loss Recovered)(均接近或达到 100%),这表明 SAE 能够有效地重构原始 ESM-2 嵌入 (embeddings),捕获其中大部分信息。同时,使用了合适的 L1 正则化参数来确保学习到的特征具有稀疏性 (sparsity)

6.3. 消融实验/参数分析

6.3.1. 稀疏性与重建质量

Table 1 中的 % Loss Recovered 指标直接反映了 SAE 的重建质量 (reconstruction quality)。高达 98% 到 100% 的损失恢复表明 SAE 能够非常精确地重构 ESM-2 的原始嵌入 (embeddings),这意味着 SAE 成功地在保持信息完整性的同时,将嵌入分解为稀疏特征。L1 正则化参数在 0.08 到 0.1 之间,这是在稀疏性 (sparsity)重建准确性 (reconstruction accuracy) 之间取得平衡的关键。

6.3.2. 随机基线的重要性

通过与在随机 ESM-2 权重上训练的 SAE 进行比较,作者验证了 SAE 捕获的生物学概念并非偶然,而是来自 PLM 预训练所学习到的有意义的表示 (representations)。Figure 3 和 Figure 5 中的绿色条清晰地表明,随机模型无法提取出与已知生物学概念强关联的特征,这有力地支持了 InterPLM 框架所提取特征的生物学有效性。

6.3.3. LLM 描述的准确性

LLM 生成的特征描述与实际特征激活的皮尔逊相关系数 (Pearson r correlation) 中位数为 0.72(Figure 8b),这表明 LLM 能够相当准确地捕捉特征的激活模式。这对于自动化特征解释 (feature interpretation) 至关重要,尤其是在处理大量特征时。此外,作者还指出 LLM 描述的准确性与 Swiss-Prot 概念预测性能的相关性较低 (Pearson r = 0.11,见 Figure 15),这意味着 LLM 能够独立地、有效地解释特征,即使它们不易被现有注释分类。

下图(原文 Figure 15)展示了 Swiss-Prot 标签质量与 LLM 描述之间的关联:

该图像是一个散点图,展示了机器学习模型的激活预测相关性与最佳Swiss-Prot概念F1值之间的关系。散点的颜色表示不同特征,回归线则显示出整体趋势。 图15:Swiss-Prot 标签质量与 LLM 描述之间存在小幅相关性 (r = 0.11)。

7. 总结与思考

7.1. 结论总结

本文成功提出了 InterPLM 框架,一种利用稀疏自编码器 (SAE)蛋白质语言模型 (PLM) 的中间层嵌入 (embeddings) 中系统性地提取人类可解释 (human-interpretable) 特征的方法。主要结论包括:

  • 克服叠加问题: SAE 有效解决了 PLM 神经元的叠加 (superposition) 现象,从 ESM-2 中解耦出多达 2,548 个比单个神经元更清晰、更具体的生物学概念特征。

  • 丰富的生物学关联: 这些 SAE 特征与多达 143 个已知 Swiss-Prot 概念 (Swiss-Prot concepts) 强关联,涵盖结合位点、结构基序、功能域等,且在定量评估 (quantitative evaluation) 中显著优于原始 ESM-2 神经元。

  • 发现新颖概念: PLM 学习到了超出现有生物学注释的连贯概念 (coherent concepts),InterPLM 框架能够识别并解释这些概念。

  • 自动化解释能力: 创新性地利用大语言模型 (LLM) 实现了对潜在特征的自动、准确的自然语言描述,大大提升了可扩展性。

  • 实际应用价值: 演示了 SAE 特征在填充蛋白质数据库中缺失注释 (filling missing annotations)定向引导蛋白质序列生成 (targeted steering of protein sequence generation) 方面的强大潜力。

  • 社区贡献: 发布了 InterPLM 交互式可视化平台 (interPLM.ai) 和相关代码,为社区提供了探索 PLM 内部机制的工具。

    总而言之,本文揭示了 PLM 内部编码了丰富而可解释的蛋白质生物学表示,并提供了一个系统框架来提取、分析和应用这些潜在特征,为蛋白质研究和模型开发开辟了新途径。

7.2. 局限性与未来工作

作者指出了以下局限性及未来研究方向:

  • 扩展到结构预测模型: 目前的工作主要集中在序列 PLM (ESM-2)。未来关键一步是将这种方法扩展到 结构预测模型 (structure-prediction models),如 ESMFoldAlphaFold。理解这些模型中学习到的特征如何演变,以及如何通过定向引导 (targeted steering) 这些特征来控制生成的蛋白质构象,将是重要的研究方向。
  • 更复杂的生物模式引导: 虽然已证明可以引导简单的周期性甘氨酸模式,但引导更复杂的生物模式仍然具有挑战性。需要开发更强大的方法和明确的成功指标来评估这种复杂模式的引导。
  • SAE 训练和评估的改进: SAE 训练方法和特征可解释性 (feature interpretability) 的评估可以通过最近 SAE 训练方法的进展和额外的验证度量 (validation metrics) 来进一步改进。
  • 未掩码氨基酸之外的信息: 目前的分析主要关注未掩码的氨基酸的特征。PLM 在掩码词元 (masked tokens) 嵌入和 <CLS><CLS> (分类 token) 嵌入中编码的信息仍待探索,这些信息对蛋白质级别的表示可能非常有价值。
  • 特征组合的理解: SAE 特征揭示了学习到的模式,但这些特征如何组合成更高级的生物学概念仍然是一个未解之谜。映射这些特征的组合方式可以揭示 PLM 内部的计算电路 (computational circuits),从而促进更深入的生物学发现。

7.3. 个人启发与批判

7.3.1. 个人启发

这篇论文提供了一个非常深刻的启发,即通过精巧的可解释性 (interpretability) 工具,我们不仅能理解“黑箱”模型的工作原理,甚至能反过来利用模型发现新的科学知识。SAE 在解耦 (disentangling) 叠加概念上的能力,是理解复杂神经网络的关键一步,它将原本模糊的神经元激活转化为原子化的、可理解的生物学实体。

  • 模型作为科学发现工具: InterPLM 证明了 PLM 不仅仅是预测工具,更是潜在的科学发现工具 (scientific discovery tool)。模型学习到的连贯概念 (coherent concepts) 提示我们,AI 可能会在现有知识体系之外发现新的生物学原理或基序,这超越了传统归纳推理 (inductive reasoning) 的范畴。
  • 可解释性是双刃剑: 可解释性不仅是为了建立信任,更是为了控制 (control)。论文中定向引导蛋白质生成 (targeted steering of protein generation) 的例子,虽然仍是初步,但预示着我们可以通过理解模型的内部特征,实现对生成式模型的精细化控制 (fine-grained control)。这对于蛋白质设计 (protein design)药物发现 (drug discovery) 具有颠覆性的潜力。
  • 跨领域方法迁移: 将 NLP 领域中用于解释 LLM 的 SAE 技术成功迁移到生物信息学领域的 PLM,体现了跨领域方法迁移 (cross-domain method transfer) 的巨大价值。这提示我们,许多在 AI 领域取得成功的方法,其核心思想可能具有普适性,值得在不同领域进行探索。
  • LLM 辅助科学研究: LLM 在自动生成特征描述中的应用是一个亮点。它表明,LLM 不仅可以处理自然语言任务,还可以作为科学研究助理 (scientific research assistant),自动化耗时且需要领域知识的解释任务,极大地提高了研究效率。

7.3.2. 批判与潜在改进

尽管 InterPLM 框架非常出色,但仍有一些潜在问题或可以改进的地方:

  • “人类可解释性”的定义和量化: 论文虽然强调了“人类可解释”特征,但在实际评估中,主要依赖于与现有 Swiss-Prot 概念的关联以及 LLM 描述的准确性。这是否真正等同于人类生物学家的直观理解,仍需更直接的人机交互实验 (human-in-the-loop experiments) 进行验证。LLM 生成的描述虽然与激活相关,但其生物学深度和新颖性是否足够,也值得进一步探讨。
  • 特征粒度的平衡: SAE 旨在解耦概念,但特征的粒度是一个权衡。过细的特征可能过于碎片化,难以形成有意义的整体概念;过粗的特征又可能无法完全解耦叠加 (superposition)。如何优化 SAE 字典大小和稀疏性参数,以找到最佳的特征粒度,是一个持续的挑战。
  • 因果关系而非相关性: 论文主要展示了特征激活与生物学概念之间的相关性 (correlation)。虽然 LLM 描述和引导实验暗示了潜在的因果关系 (causal relationship),但更严格的因果分析(例如通过干预研究 (interventional studies))将更有力地证明这些特征是模型做出特定预测的真正驱动因素。
  • 复杂相互作用的缺失: 蛋白质的生物学功能往往涉及多个基序和域的复杂协同作用 (cooperative interactions)。目前的 SAE 框架主要提取原子化的特征,如何将这些原子特征组合起来,解释更高级别的相互作用网络 (interaction networks)功能回路 (functional circuits),是未来需要深入探索的方向。
  • 计算成本和可扩展性: 训练多个 SAE 来解释 PLM 的每一层,尤其是对于更大的 PLM 模型,可能会带来显著的计算成本。如何提高 SAE 训练的效率和可扩展性,使其能够应用于前沿的超大规模 PLM,是一个实际挑战。
  • 可视化平台的进一步优化: InterPLM 平台是一个很好的社区资源。未来可以增加更多交互式分析工具,例如允许用户自定义概念搜索、特征组合分析、或者与外部结构可视化工具的深度集成,以进一步赋能生物学家。

相似论文推荐

基于向量语义检索推荐的相关论文。

暂时没有找到相似论文。