SaProt: 融合结构字母表(AA+3Di)的蛋白质语言模型¶
概述¶
SaProt 是一种结构与序列协同感知的蛋白质语言模型(Protein Language Model, PLM)。其核心创新在于提出了 结构感知词汇表(structure-aware vocabulary) ——将经典的氨基酸(Amino Acid, AA)与折叠描述符 3Di token 拼合,形成“AA+3Di”的复合token。这一设计使模型在预训练和下游任务中能够同时捕获一维序列信息与三维结构约束,显著提升了突变效应预测、功能注释、稳定性评估等任务的性能。
本模型已发表于 Nature Biotechnology,并在 ICLR 2024 获得 Spotlight。SaProt 提供 35M、650M 和 1.3B 参数版本,均开源在 Hugging Face。最新发布的 ColabSaprot v2 和 SaprotHub 平台进一步降低了生物学家使用大模型的技术门槛,且在湿实验验证中取得成功。
关键特性: - 以“AA+3Di”字母表编码蛋白质,输入长度翻倍但保留完整序列顺序 - 基于大量 AlphaFold2 预测结构与 PDB 实验结构预训练 - 全面超越 ESM-2 等同规模模型,尤其在需要结构感知的任务上 - 支持零样本突变效应预测、嵌入提取、逆折叠设计
核心知识点¶
1. 结构感知序列 (SA token)¶
SaProt 将每一个残基表示为一个 结构感知 token (SA token) ,由两部分拼接: - AA: 标准 20 种氨基酸 - 3Di: 由 Foldseek 生成的结构字母表,离散为20种三维相互作用状态
例如,一个 100 个残基的蛋白质,传统序列表示为 [L, A, G, ...],而 SaProt 表示为 [L_A, A_B, G_C, ...],其中后半部分为对应残基的 3Di token。词汇表大小从 30 左右(含特殊token)扩展为 20 × 20 = 400 种组合 token,形成结构感知的排序输入。
2. 模型版本与规模¶
| 模型标识 | 参数量 | 预训练数据 | 结构输入要求 |
|---|---|---|---|
| SaProt_35M_AF2 | 35M | 40M AF2 预测结构 | 需要 SA token(也提供纯序列版本用于对比) |
| SaProt_650M_AF2 | 650M | 40M AF2 结构 | 需要 SA token,纯序列模式微调后方可工作 |
| SaProt_650M_PDB | 650M | 40M AF2 + 60K PDB 结构 | 同上 |
| SaProt_1.3B_AF2 | 1.3B | 40M AF2 结构 | 支持 SA token 和纯 AA 序列,两者皆可 |
| SaProt_1.3B_AFDB_OMG_NCBI | 1.3B | 40M AF2 + 200M OMG + 150M NCBI | 同上,序列模式表现尤佳 |
⚠️ 重要提示:
- 35M 和 650M 版本在不提供结构 token 时性能骤降,其 frozen embeddings 仅对 SA 输入有效。若使用纯序列,需进行微调。
- 1.3B 版本原生支持高质量 AA 序列输入,但在有高质量结构时,SA 模式仍优于纯序列模式,且预期可超过 ESM-2。
3. 性能亮点¶
SaProt 在多个蛋白质任务基准上显著超越同尺寸 ESM-2,尤其在需要结构理解的场景中优势明显。以 650M 模型为例:
| 任务 | 指标 | ESM-2 650M | SaProt 650M |
|---|---|---|---|
| ClinVar 致病性预测 | AUC | 0.880 | 0.911 |
| ProteinGym 突变效应 | Spearman ρ | 0.471 | 0.520 |
| Thermostability 热稳定性 | Spearman ρ | 0.738 | 0.769 |
| EC 酶学功能预测 | Fmax | 0.874 | 0.902 |
| GO-MF 分子功能 | Fmax | 0.665 | 0.702 |
在 ProteinGym 公开排行榜上,SaProt 曾位列第一,且是唯一以单一通用结构模型进入前列的方法,其余均为混合模型或突变特化模型。
4. 架构与预训练策略¶
SaProt 沿用 ESM-2 的 Transformer 骨架,仅将输入嵌入层从 (vocab_size_aa, dim) 替换为 (vocab_size_SA, dim),其余结构保持不变。预训练任务为标准的掩码语言建模(Masked Language Modeling, MLM),使用 40M 来自 AlphaFold2 的结构预测数据,部分版本加入 PDB 实验结构进行第二阶段训练。1.3B 版本在 AF2 基础上融入大规模宏基因组序列(OMG、NCBI),使用 70% 序列一致性过滤。
代码实操¶
环境安装¶
# 创建虚拟环境
conda create -n SaProt python=3.10
conda activate SaProt
# 一键安装依赖(PyTorch, fair-esm, foldseek等)
bash environment.sh
模型加载 (Hugging Face)¶
from transformers import AutoTokenizer, AutoModelForMaskedLM
tokenizer = AutoTokenizer.from_pretrained("westlake-repl/SaProt_650M_AF2")
model = AutoModelForMaskedLM.from_pretrained("westlake-repl/SaProt_650M_AF2")
# 准备结构感知序列 (AA+3Di)
sa_sequence = ["M_A", "E_B", "T_C", ...] # 需先通过结构转换获得
inputs = tokenizer(sa_sequence, return_tensors="pt", is_split_into_words=True)
outputs = model(**inputs)
蛋白质结构转为 SA 序列¶
需要借助 Foldseek 从 PDB/mmCIF 结构文件中生成 3Di token。
# 将结构文件转为 3Di 序列
foldseek createdb example.pdb tmp_db
foldseek lndb tmp_db_h tmp_db_ss_h
foldseek convert2fasta tmp_db_ss tmp_db_ss.fasta
foldseek createindex tmp_db tmp_output
foldseek search tmp_db tmp_db tmp_result tmp --threads 8
foldseek aln2tmscore tmp_db tmp_db tmp_result tmp_tmscore
foldseek unpackdb tmp_db_ss tmp_db_ss.fasta --output-fmt fasta
在 Python 中处理生成的 3Di 序列并与 AA 序列合并。
from saprot.utils.foldseek_util import get_struc_seq
# 指定 Foldseek 路径和结构文件
foldseek_path = "/path/to/foldseek"
pdb_path = "example.pdb"
parsed = get_struc_seq(foldseek_path, pdb_path, chains=["A"]) # 返回 {"A": ("AA seq", "3Di seq")} # 中文注释:返回链的氨基酸和3Di序列
aa_seq, foldseek_seq = parsed["A"]
sa_seq = [f"{aa}_{foldseek}" for aa, foldseek in zip(aa_seq, foldseek_seq)] # 中文注释:拼合成SA token列表
零样本突变效应预测¶
SaProt 支持无需额外训练直接预测点突变对蛋白质功能的影响(打分基于掩码位置的对数概率比)。
from saprot.utils.mutation import predict_mutation_effect
wt_sa_seq = "M_A E_B T_C D_D ...".split()
mut_pos = 2 # 0-indexed,突变位置
mut_aa = "A" # 突变氨基酸
# 中文注释:计算突变效应得分,正值代表突变可能有益
score = predict_mutation_effect(model, tokenizer, wt_sa_seq, mut_pos, mut_aa)
print(f"Mutation effect score: {score:.3f}")
提取蛋白质嵌入¶
# 获取整条序列的表示
from saprot.utils.embedding import get_protein_embedding
embedding = get_protein_embedding(model, tokenizer, sa_sequence, layer=-1, pooling="mean") # 中文注释:取最后一层均值池化
print(embedding.shape) # e.g., (1280,) for 650M model
蛋白质逆折叠¶
SaProt 可用于从结构(3Di)恢复氨基酸序列,即反向预测最可能的 AA token。
from saprot.utils.inverse_folding import inverse_folding
sa_seq_mask_aa = ["<mask>_A", "<mask>_B", ...] # 中文注释:仅遮蔽AA部分,保留3Di
recovered_aa = inverse_folding(model, tokenizer, sa_seq_mask_aa)
print("".join(recovered_aa))
常见问题¶
Q1: 为什么 SaProt 35M/650M 用纯氨基酸序列效果很差?
A: 这些版本在预训练时只见过 SA token 输入,模型嵌入层和位置编码均针对 400 个 token 学习。直接输入 AA token 会导致分布不匹配,必需经过微调才能迁移。1.3B 版本在预训练阶段同时暴露了纯序列与 SA 序列,因此可以泛化。
Q2: 我只有蛋白质序列,没有结构,能否使用 SaProt?
推荐使用 SaProt_1.3B_AF2 或 SaProt_1.3B_AFDB_OMG_NCBI,它们在 AA 序列模式下表现良好。如果仍希望使用 650M 版本,可用 ESMFold 或 AlphaFold2 在线生成结构,然后生成 3Di token 输入。
Q3: 如何获取 3Di token?
安装 Foldseek,将 PDB/mmCIF 结构转换为 3Di 序列(上述代码已包含)。ColabSaprot 界面则内置了自动化流程。
Q4: SaProt 在预测突变效应时是否需要结构信息?
零样本预测通常需要结构感知序列,但 1.3B 序列模式也能工作,结构会带来进一步提升。
Q5: 微调 SaProt 需要什么数据?
对于监督任务(如功能分类),数据集应按 30% 序列一致性划分,以避免同源性泄露。SaProt 论文中各类任务的分割方式均遵循此原则。
Q6: ColabSaprot 与 SaprotHub 是什么?
ColabSaprot 是一个无需编程、浏览器内运行的蛋白质分析与设计工具,集成了 SaProt 模型。SaprotHub 则提供了训练蛋白质语言模型的框架,允许生物学家用自己的数据微调模型。两者均已得到实验验证。
速查表¶
模型下载资源¶
| 模型名称 | Hugging Face 地址 | 参数量 | 适用性 |
|---|---|---|---|
| SaProt_35M_AF2 | westlake-repl/SaProt_35M_AF2 | 35M | 轻量级,需要结构 token |
| SaProt_650M_AF2 | westlake-repl/SaProt_650M_AF2 | 650M | 强结构依赖,推荐首选 |
| SaProt_650M_PDB | westlake-repl/SaProt_650M_PDB | 650M | 额外 PDB 训练 |
| SaProt_1.3B_AF2 | westlake-repl/SaProt_1.3B_AF2 | 1.3B | 序列/结构双模式,性能最佳 |
| SaProt_1.3B_AFDB_OMG_NCBI | westlake-repl/SaProt_1.3B_AFDB_OMG_NCBI | 1.3B | 更多预训练数据,泛化性更强 |
关键函数速览¶
| 功能 | 模块/函数 | 输入 | 输出 |
|---|---|---|---|
| 结构转 SA 序列 | saprot.utils.foldseek_util.get_struc_seq | PDB 路径 | (AA 序列, 3Di 序列) |
| 零样本突变预测 | saprot.utils.mutation.predict_mutation_effect | SA 序列+突变信息 | 分数 |
| 蛋白质嵌入提取 | saprot.utils.embedding.get_protein_embedding | SA 序列 | 嵌入向量 |
| 逆折叠 | saprot.utils.inverse_folding.inverse_folding | 结构感知序列(AA 掩码) | 预测氨基酸序列 |
任务性能速查(650M,部分指标)¶
- ClinVar AUC: 0.911
- ProteinGym Spearman ρ: 0.520
- Thermostability Spearman ρ: 0.769
- HumanPPI Acc: 83.50%
- EC Fmax: 0.902
- GO-MF Fmax: 0.702
- Subcellular Acc: 84.22%
以上结果均基于 30% 序列一致性分割的监督微调或零样本设置,详细数据参见原始论文。
扩展阅读与引用: - 论文: SaProt: Protein Language Modeling with Structure-aware Vocabulary (bioRxiv 2024, ICLR 2024, Nat Biotechnol 2025) - 官方仓库: westlake-repl/SaProt - SaprotHub 训练平台: SaprotHub - ColabSaprot 在线工具: ColabSaprot v2