蛋白质大模型新突破：SaProt通过结构感知词汇提升预测精度
蛋白质大模型新突破：SaProt通过结构感知词汇提升预测精度
蛋白质语言模型（PLM）在生物信息学领域发挥着重要作用，但传统的PLM往往缺乏对蛋白质结构信息的明确考虑。为了解决这一问题，研究者们提出了具有结构感知词汇的蛋白质语言模型SaProt。该模型通过引入结构感知的词汇表，将蛋白质的结构信息编码为离散的3D tokens，从而构建了一个结构感知的序列。这种创新的方法使得模型能够在大规模预训练中学习到蛋白质结构与序列之间的关系。
SAPROT：具有结构感知词汇的蛋白质语言模型
论文地址：https://www.biorxiv.org/content/10.1101/2023.10.01.560349v1.full.pdf
项目地址：https://github.com/westlake-repl/SaProt
模型地址：https://huggingface.co/westlake-repl/SaProt_650M_AF2
研究背景及项目介绍：
大规模蛋白质语言模型（PLM），例如ESM系列，通过对残基序列进行无监督训练，在与蛋白质结构和功能相关的各种下游任务中取得了显着的性能。 它们已成为生物学研究人员和实践者的重要工具。 然而，普通 PLM 的一个局限性是缺乏对蛋白质结构信息的明确考虑，这表明有进一步改进的潜力。 受此启发，我们引入了“结构感知词汇”的概念。
SaProt 是一个创新的蛋白质语言模型，它通过引入结构感知的词汇表，将蛋白质的结构信息编码为离散的3D tokens，从而构建了一个结构感知的序列。这种创新的方法使得模型能够在大规模预训练中学习到蛋白质结构与序列之间的关系。SaProt不仅在多个任务上超越了现有的模型，如ESM-2，还在ProteinGym基准测试中取得了第一名的成绩（2024年10月）。
参考网站：https://proteingym.org/benchmarks
具有结构感知词汇表的 SaProt 框架
利用 Foldseek 工具对蛋白质结构进行编码，将其离散化为 3D tokens，每个 3D token 代表一个残基的几何构象信息。结合残基（例如图中的“ACHY”）和结构 tokens，创建结构感知序列，形成 “SA” 词汇。
模型输入：结构感知序列、掩码处理后的残基序列。
模型输出：1.预测的残基序列token 2.针对不同任务给出对应结果（如在 ProteinGym 和 ClinVar 数据集中，分别输出与评估指标（Spearman’s ρ 和 AUC）相关的预测值。）
模型架构：SaProt采用与ESM2的650M版本相同的网络架构和参数大小。 主要区别在于扩展的嵌入层，它包含 441 个 SA 令牌，而不是原来的 20 个剩余令牌。 这种几乎相同的架构可以与 ESM 模型进行直接比较。
目标函数：使用类似于 ESM-1b 和 ESM-2 的 BERT 式 MLM 目标来训练 SaProt，从而支持蛋白质级别和残基级别任务。 掩码策略1：最直接的掩蔽策略是使用符号“##”随机掩蔽几个SA令牌sifi，然后直接从SA词汇表中预测它们。 但如果 SA 令牌不够准确，预测准确的 SA 令牌可能会导致模型走向错误的优化方向。 掩码策略2：另一种潜在的掩蔽策略涉及预测残差标记 si 或预测 Foldseek 结构标记 fi。 然而，预测 fi 遇到了上面提到的同样的问题。 由于为了提高蛋白质一级序列中残基类型的准确性，仅预测残基标记似乎是更有效的训练方法。
实验结果
1.零样本突变效应预测任务 数据集与评估指标：
采用 ProteinGym 和 ClinVar 数据集，分别使用 Spearman’s ρ 和 AUC 作为评估指标。
实验结果： SaProt 在该任务上优于所有基于残基序列和结构的基线模型，包括著名的 ESM 系列模型。与 ESM - 2 相比，SaProt 在预测突变效应方面有持续改进；与 MIF - ST 相比，尽管后者使用实验确定的高精度结构训练，但 SaProt 由于大规模结构训练数据表现更优。
多序列比对（MSA）信息可增强模型零样本能力，将该技术应用于 SaProt 和基线模型后，SaProt 仍保持最高准确率，表明适用于残基序列模型的改进技术对 SaProt 也可能有用。 ESM系列由 Meta开发，是基于深度学习的蛋白质语言模型，通过对大量蛋白质序列进行无监督学习，捕捉蛋白质序列中的信息，用于预测蛋白质的各种性质和功能。ESM-2 是目前规模最大、最复杂的蛋白质语言模型之一，参数量达 150 亿。ESMFold 在蛋白质结构预测的准确性上与 AlphaFold2相当，并且速度更快。
2.监督微调任务
数据集与任务类型：在多个数据集上评估 SaProt 在蛋白质级任务上的性能，包括预测热稳定性、金属离子结合、蛋白质定位（DeepLoc）、蛋白质注释（EC 和 GO）以及蛋白质 - 蛋白质相互作用（HumanPPI）等任务。
结果分析：SaProt 在所有蛋白质级任务中均优于 ESM - 2，在热稳定性、HumanPPI、金属离子结合和 DeepLoc 等任务中表现出显著增强，再次证明整合结构信息可提升蛋白质表示能力。 SaProt 大幅优于 GearNet 和 MIF - ST 等结构模型，凸显其结构建模的有效性。
SaProt - GearNet 组合模型在某些情况下优于 ESM - GearNet，表明 SaProt 与其他先进改进技术具有正交性，但组合模型不一定总是优于单个模型。
验证加入结构序列的可行性
1.蛋白质结构感知能力
通过在 TAPE 基准上进行接触图预测任务（冻结骨干网络，仅微调接触头），评估 SaProt 对蛋白质结构信息的掌握程度。结果表明，SaProt 在接触图预测任务中显著优于 ESM - 2，表明其包含更准确的结构信息，从而预期在蛋白质功能预测任务中具有更高准确性。
进行额外的零样本预测实验，随机替换 SaProt 中的部分结构 token 为随机（Foldseek）token，发现随着替换比例增加，SaProt 的性能下降，强调了结构 token 对模型性能的重要性。
2.PDB 与 AlphaFoldDB 结构对模型性能的影响
研究了在实验确定结构（PDB 结构）上训练的 SaProt - PDB 模型与仅在 AF2 结构上训练的 SaProt 模型的性能差异。在下游任务评估中，发现仅在 AF2 结构上训练时，SaProt 在 AF2 和 PDB 结构上的整体准确性差异不大，但对于 SaProt - PDB 模型，在有 PDB 结构可用时直接使用可获得更好的准确性，尤其在零样本任务中，训练和测试数据一致性对性能有重要影响。
3.模型可视化分析
使用 t - SNE 可视化 SaProt 和 ESM - 2 在 SCOPe 数据库上生成的蛋白质表示，发现 SaProt 能更好地区分不同结构类型的蛋白质，如 alpha 和 beta 蛋白，进一步证明其对结构变化的辨别能力。
可视化所有 400 个结构感知 token 的嵌入，观察到一定程度的聚类现象，语义空间中相近的 SA token 通常对应相似类型的残基或 Foldseek token。
结论
提出了一种结构感知（SA）词汇，将蛋白质的一级和三级结构信息整合到 SA - token 中，以此为基础训练的通用 PLM SaProt 在 10 个蛋白质功能预测任务中取得了最先进的性能，为生物学研究提供了有力工具。
局限性分析
模型性能依赖问题：SA 词汇的性能依赖于 Foldseek，其在平衡搜索效率和编码准确性方面存在改进空间，从而限制了 SaProt 的表示能力。
模型规模受限：由于计算资源限制，SaProt 的模型规模可能未达最大潜力。
潜在应用探索：除已研究任务外，SA 词汇在其他应用中具有潜力，如预测蛋白质复合物结构（通过将两个蛋白质序列替换为 SA - token 序列，自动整合单链结构信息）和蛋白质生成任务（生成 SA - token 序列可能在生成过程中提供更强的结构约束），这些方向有待未来研究探索。