esm-protein-language-model

ESM — 蛋白质语言模型

概览

ESM（Evolutionary Scale Modeling）提供预训练的蛋白质语言模型，用于生成式蛋白质设计和表示学习。ESM3 是一种多模态生成模型，可同时根据序列、结构和功能进行条件生成。ESM C 是一种高效的嵌入模型，经过优化，可提取蛋白质表示，用于下游 ML 任务。

适用场景

• 根据所需的结构或功能生成新型蛋白质序列
• 从蛋白质序列中提取固定长度的嵌入，用于分类、聚类或回归
• 根据氨基酸序列预测 3D 结构
• 反向折叠：设计能够折叠成目标结构的序列
• 使用功能关键词（GO 术语、EC 编号）注释蛋白质
• 通过嵌入距离而非序列比对比较蛋白质相似性
• 思维链蛋白质设计：迭代优化序列／结构／功能
• 对于传统的基于物理的结构预测，请改用 AlphaFold
• 对于序列比对和同源性搜索，请通过 BioPython 改用 BLAST/HMMER

前提条件

• Python 包：esm（EvolutionaryScale 包）
• 硬件：建议使用 GPU 进行本地推理（ESM3：8GB+ VRAM；ESM C：4GB+ VRAM）。小批次也可使用 CPU
• 云端替代方案：EvolutionaryScale Forge API（需要从 forge.evolutionaryscale.ai 获取 API 令牌）
• 模型权重：首次使用时自动下载（约 1-4 GB，具体取决于模型）

pip install esm
# For Forge cloud API
pip install esm[forge]

快速开始

from esm.models.esmc import ESMC
from esm.sdk.api import ESMProtein

# Load ESM C model for embeddings
model = ESMC.from_pretrained("esmc_600m")

# Create protein from sequence
protein = ESMProtein(sequence="MKTAYIAKQRQISFVKSHFSRQLEERLGLIEVQAPILSRVGDGTQDNLSGAEKAVQVKVKALPDAQFEVVHSLAKWKRQQIAATGFHIIPGDKPDNRAGGYDN")

# Get per-residue embeddings
output = model(protein)
embeddings = output.embeddings  # shape: (1, seq_len, embedding_dim)
print(f"Embedding shape: {embeddings.shape}")
# Embedding shape: (1, 101, 1152)

核心 API

1. 蛋白质序列生成（ESM3）

根据结构、功能或部分序列生成新型蛋白质序列。

from esm.models.esm3 import ESM3
from esm.sdk.api import ESM3InferenceClient, ESMProtein, GenerationConfig

# Load ESM3 locally
model = ESM3.from_pretrained("esm3_sm_open_v1")

# Generate from partial sequence (fill in masked positions)
prompt = ESMProtein(sequence="MKTAYIAK____ISFVK____RQLEERLG")  # ____ = positions to generate
config = GenerationConfig(track="sequence", num_steps=10, temperature=0.7)
generated = model.generate(prompt, config)
print(f"Generated sequence: {generated.sequence[:50]}...")

# Conditional generation: design sequence for a target structure
from esm.sdk.api import ESMProtein, GenerationConfig
from esm.utils.structure.protein_chain import ProteinChain

# Load target structure from PDB
chain = ProteinChain.from_pdb("target.pdb")
prompt = ESMProtein.from_protein_chain(chain)
prompt.sequence = None  # Clear sequence, keep structure

config = GenerationConfig(track="sequence", num_steps=16, temperature=0.5)
designed = model.generate(prompt, config)
print(f"Designed sequence ({len(designed.sequence)} residues): {designed.sequence[:50]}...")

2. 蛋白质嵌入（ESM C）

提取固定长度的表示，用于下游 ML 任务。

from esm.models.esmc import ESMC
from esm.sdk.api import ESMProtein
import torch

model = ESMC.from_pretrained("esmc_600m")  # or "esmc_300m" for lighter model

sequences = [
    "MKTAYIAKQRQISFVKSHFSRQLEERLGLIEVQAPILSRVGDGTQDNLSGAEKAVQVKVKALPDAQFEVVHSLAKWKRQQIAATGFHIIPGDKPDNRAGGYDN",
    "MKWVTFISLLFLFSSAYSRGVFRRDAHKSEVAHRFKDLGEENFKALVLIAFAQYLQQCPFEDHVKLVNEVTEFAKTCVADESAENCDKS",
]

embeddings = []
for seq in sequences:
    protein = ESMProtein(sequence=seq)
    output = model(protein)
    # Mean-pool per-residue embeddings to get fixed-length vector
    mean_emb = output.embeddings.mean(dim=1)  # shape: (1, embedding_dim)
    embeddings.append(mean_emb)

emb_matrix = torch.cat(embeddings, dim=0)
print(f"Embedding matrix: {emb_matrix.shape}")  # (2, 1152)

# Compute pairwise similarity
similarity = torch.cosine_similarity(emb_matrix[0:1], emb_matrix[1:2])
print(f"Cosine similarity: {similarity.item():.4f}")

3. 结构预测

根据氨基酸序列预测 3D 坐标。

from esm.models.esm3 import ESM3
from esm.sdk.api import ESMProtein, GenerationConfig

model = ESM3.from_pretrained("esm3_sm_open_v1")

protein = ESMProtein(sequence="MKTAYIAKQRQISFVKSHFSRQLEERLGLIEVQAPILSRVGDGTQDNLSGAEKAVQVKVKALPDAQFEVVHSLAKWKRQQIAATGFHIIPGDKPDNRAGGYDN")

# Generate structure from sequence
config = GenerationConfig(track="structure", num_steps=16)
result = model.generate(protein, config)

# Save predicted structure
result.to_pdb("predicted.pdb")
print(f"Saved structure: {len(result.sequence)} residues → predicted.pdb")

4. 反向折叠

设计能够折叠成目标 3D 结构的氨基酸序列。

from esm.models.esm3 import ESM3
from esm.sdk.api import ESMProtein, GenerationConfig
from esm.utils.structure.protein_chain import ProteinChain

model = ESM3.from_pretrained("esm3_sm_open_v1")

# Load target structure
chain = ProteinChain.from_pdb("target_structure.pdb")
prompt = ESMProtein.from_protein_chain(chain)

# Clear sequence but keep structure coordinates
prompt.sequence = None

# Generate multiple designs
designs = []
for i in range(5):
    config = GenerationConfig(track="sequence", num_steps=16, temperature=0.7)
    designed = model.generate(prompt, config)
    designs.append(designed.sequence)
    print(f"Design {i+1}: {designed.sequence[:40]}...")

print(f"Generated {len(designs)} sequence designs for target structure")

5. 功能条件生成

生成具有所需功能注释（GO 术语、酶活性）的蛋白质。

from esm.models.esm3 import ESM3
from esm.sdk.api import ESMProtein, GenerationConfig

model = ESM3.from_pretrained("esm3_sm_open_v1")

# Condition on functional keywords
protein = ESMProtein(
    sequence=None,  # generate de novo
    function_annotations=["ATP binding", "kinase activity", "protein phosphorylation"],
)

config = GenerationConfig(track="sequence", num_steps=32, temperature=0.7)
result = model.generate(protein, config)
print(f"Function-conditioned sequence: {result.sequence[:50]}...")
print(f"Length: {len(result.sequence)} residues")

6. Forge 云端 API

使用 EvolutionaryScale 的云端推理能力运行大型模型，无需本地 GPU。

from esm.sdk.forge import ESM3ForgeInferenceClient
from esm.sdk.api import ESMProtein, GenerationConfig

# Authenticate (requires FORGE_API_TOKEN env var or explicit token)
client = ESM3ForgeInferenceClient(model="esm3-open-2024-03", token="your_token_here")

protein = ESMProtein(sequence="MKTAYIAKQRQISFVKSHFSRQLEERLG")
config = GenerationConfig(track="structure", num_steps=16)
result = client.generate(protein, config)

result.to_pdb("forge_predicted.pdb")
print("Predicted structure via Forge API → forge_predicted.pdb")

核心概念

ESM3 与 ESM C：如何选择

特性	ESM3	ESM C
主要用途	生成式蛋白质设计	嵌入提取
能力	序列生成、结构预测、反向折叠、功能条件生成	逐残基和均值池化嵌入
模型大小	esm3_sm_open_v1（约 1.4B 个参数）	esmc_300m、esmc_600m
GPU 要求	8GB+ VRAM	4GB+ VRAM（esmc_300m：2GB）
使用场景	设计新蛋白质、预测结构	下游 ML（分类、聚类、回归）
云端选项	Forge API（可用更大的模型）	仅限本地

GenerationConfig 参数

GenerationConfig 控制 ESM3 生成输出的方式：

• track：要生成的模态（"sequence"、"structure"、"function"）
• num_steps：迭代优化步数（越多 = 质量越高、速度越慢）
• temperature：采样温度（0.0 = 贪心采样，0.5-0.7 = 多样化，1.0 = 最大多样性）

ESMProtein 对象

用于存放序列、结构坐标和功能注释的核心数据容器：

• .sequence — 氨基酸字符串（例如："MKTAY..."）
• .coordinates — 原子 3D 位置（Nx3 张量）
• .function_annotations — 功能关键词列表
• 使用 ESMProtein.from_protein_chain() 从 PDB 结构加载
• 使用 .to_pdb() 保存预测的结构

常见工作流

工作流 1：基于蛋白质嵌入的分类

目标：从蛋白质序列中提取嵌入，并训练下游分类器。

from esm.models.esmc import ESMC
from esm.sdk.api import ESMProtein
import torch
import numpy as np

model = ESMC.from_pretrained("esmc_600m")

# Embed a set of sequences
sequences = ["MKTAY...", "MKWVT...", "MSGLI..."]  # replace with actual sequences
labels = [0, 1, 0]  # binary labels

embeddings = []
for seq in sequences:
    protein = ESMProtein(sequence=seq)
    output = model(protein)
    mean_emb = output.embeddings.mean(dim=1).detach().cpu().numpy()
    embeddings.append(mean_emb.squeeze())

X = np.array(embeddings)
y = np.array(labels)
print(f"Feature matrix: {X.shape}")  # (n_samples, 1152)

# Train a simple classifier
from sklearn.linear_model import LogisticRegression
clf = LogisticRegression(max_iter=1000).fit(X, y)
print(f"Training accuracy: {clf.score(X, y):.2f}")

工作流 2：结构条件蛋白质设计

目标：设计多个新颖序列，使其折叠成目标结构，然后按预测质量排序。

1. 使用 ProteinChain.from_pdb()（核心 API 模块 4）从 PDB 加载目标结构
2. 创建仅包含结构、不包含序列的 ESMProtein 提示
3. 使用 temperature=0.7 生成 10 个以上的序列设计，以提高多样性（核心 API 模块 1）
4. 对每个设计，根据设计序列预测结构（核心 API 模块 3）
5. 将预测结构与目标结构进行比较（计算 RMSD），为设计排序
6. 选出排名靠前的设计进行实验验证

关键参数

参数	模块/函数	默认值	范围 / 选项	效果
num_steps	GenerationConfig	varies	1–64	迭代优化步数；步数越多，质量越高、速度越慢
temperature	GenerationConfig	1.0	0.0–1.5	采样多样性；0.0 = 贪心采样，0.7 = 平衡，1.0 以上 = 更具创造性
track	GenerationConfig	—	"sequence"、"structure"、"function"	要生成的模态
模型名称	from_pretrained	—	"esm3_sm_open_v1"、"esmc_300m"、"esmc_600m"	模型大小与能力之间的权衡
token	ESM3ForgeInferenceClient	环境变量	API 令牌字符串	Forge 云端身份验证

最佳实践

1. 嵌入任务使用 ESM C，生成任务使用 ESM3：ESM C 更小、速度更快，并针对表示质量进行了优化。只有在需要生成能力（序列设计、结构预测、逆向折叠）时才使用 ESM3。

2. 对每个残基的嵌入取平均池化，获得固定长度表示：ESM C 输出每个残基的嵌入（seq_len × dim）。对于需要固定长度输入的下游机器学习任务，可沿序列维度取平均值：embeddings.mean(dim=1)。

3. 蛋白质设计使用 0.5–0.7 的温度：温度为 1.0 时会生成多样性很高、但可能不具备功能的序列。温度为 0.5–0.7 时能平衡多样性和质量。只有在确定性结构预测时才使用温度 0.0。

4. 增加 num_steps 以提高生成质量：增加迭代优化步数可提升输出质量，但会增加计算时间。快速探索使用 8–16 步，最终设计使用 32 步以上。

5. 批量处理序列以充分利用 GPU：一次处理一个序列无法充分利用 GPU。嵌入多个序列时，请将它们分批处理（批量大小受 VRAM 限制）。

6. 大规模推理或大模型推理使用 Forge API：开放权重的 ESM3 是较小的变体。如需生产级蛋白质设计，Forge API 可提供对更大模型的访问。

常用配方

配方：蛋白质两两相似度矩阵

from esm.models.esmc import ESMC
from esm.sdk.api import ESMProtein
import torch
import numpy as np

model = ESMC.from_pretrained("esmc_300m")

sequences = {
    "Protein_A": "MKTAYIAKQRQISFVK...",
    "Protein_B": "MKWVTFISLLFLFSSAYS...",
    "Protein_C": "MSGLILQRAAVIAAGASSAG...",
}

# Extract embeddings
embs = {}
for name, seq in sequences.items():
    protein = ESMProtein(sequence=seq)
    output = model(protein)
    embs[name] = output.embeddings.mean(dim=1).detach().squeeze()

# Compute similarity matrix
names = list(embs.keys())
sim_matrix = np.zeros((len(names), len(names)))
for i, n1 in enumerate(names):
    for j, n2 in enumerate(names):
        sim_matrix[i, j] = torch.cosine_similarity(embs[n1].unsqueeze(0), embs[n2].unsqueeze(0)).item()

print("Similarity matrix:")
for i, name in enumerate(names):
    print(f"  {name}: {sim_matrix[i].round(3)}")

配方：保存和加载嵌入以供重复使用

from esm.models.esmc import ESMC
from esm.sdk.api import ESMProtein
import torch
import numpy as np

model = ESMC.from_pretrained("esmc_600m")

# Generate and save
protein = ESMProtein(sequence="MKTAYIAKQRQISFVK...")
output = model(protein)
np.save("embedding.npy", output.embeddings.detach().cpu().numpy())
print("Saved embedding.npy")

# Load later (no GPU needed)
embedding = np.load("embedding.npy")
print(f"Loaded embedding: {embedding.shape}")

故障排除

问题	原因	解决方案
CUDA out of memory	模型对于 GPU 来说过大	使用更小的模型（esmc_300m）、减小批量大小，或使用 Forge 云端 API
RuntimeError: no CUDA device	没有可用的 GPU	模型可在 CPU 上运行（速度较慢）。设置 device="cpu" 或使用 Forge API
生成速度慢	num_steps 过多或使用 CPU 推理	减少 num_steps（草稿使用 8 步）、使用 GPU，或使用 Forge API 运行大模型
ImportError: esm	未安装软件包	pip install esm（注意：这是 EvolutionaryScale 的 esm，不是较早的 Facebook Research esm）
生成的序列质量低	温度过高或步数过少	将温度降低至 0.5，将 num_steps 增加至 32 步以上
Forge API 身份验证错误	API 令牌无效或缺失	设置 FORGE_API_TOKEN 环境变量，或显式传入 token=；从 forge.evolutionaryscale.ai 获取令牌
加载模型权重时出现 KeyError	模型名称错误	使用准确的名称："esm3_sm_open_v1"、"esmc_300m"、"esmc_600m"

相关技能

• alphafold-database-access — 从 AlphaFold DB 获取预测结构；使用 ESM 进行从头结构预测
• biopython — 序列 I/O 和比对；在进行 ESM 嵌入前预处理序列
• scikit-learn — 基于 ESM 嵌入进行下游 ML（分类、聚类、回归）
• rdkit — 用于小分子药物设计的化学信息学（可与蛋白质设计互补）

参考资料

• ESM GitHub — 源代码和模型权重
• EvolutionaryScale Forge — 云端推理 API
• Hayes 等人（2024）“使用语言模型模拟 5 亿年的进化” — bioRxiv
• Lin 等人（2023）“使用语言模型进行进化尺度的原子级蛋白质结构预测” — Science