ESM-2：用于生成和优化目标蛋白肽结合剂
社区文章 发布于 2023 年 11 月 23 日
在本文中，我们将讨论如何使用 ESM-2 的掩码语言建模功能，对 ESM-2 进行微调，以生成目标蛋白的肽结合剂。然后，我们将讨论如何对蛋白质-肽复合物进行计算机指导进化，以优化肽结合剂。我们还将讨论如何将其扩展到一般的相互作用蛋白质对。
引言
在快速发展的蛋白质工程领域，人工智能和机器学习技术的应用日益突出。其中，ESM-2 作为专门为蛋白质序列量身定制的先进语言模型，已显示出非凡的潜力。本文深入探讨了 ESM-2 在生成和优化靶向特定蛋白质的肽结合剂方面的实际应用。我们将探索使用掩码语言建模方法对 ESM-2 进行微调的过程，并演示如何利用它来生成肽结合剂。此外，我们还将讨论计算机指导进化技术，特别是针对蛋白质-肽复合物，以增强这些肽的结合亲和力。这项探索不仅突出了 ESM-2 在蛋白质工程中的潜力，也为它在理解和设计蛋白质-蛋白质相互作用方面更广泛的应用奠定了基础。
我们将使用模型 PepMLM，它是 ESM-2 的微调版本，但您可以通过克隆 PepMLM 的 Github 仓库轻松微调您自己的 ESM-2 版本。前往Github 仓库并克隆它！接下来，您只需要调整训练数据和测试数据的文件路径，然后运行训练脚本。该模型在普通的 GPU 上训练速度相当快。您也可以尝试使用 UniProt 中可用的蛋白质-蛋白质相互作用数据，在您自己的数据集上微调 ESM-2。如果您的相互作用数据集更大，这可能会产生更好的结果。
使用 PepMLM 生成肽结合剂
首先，我们需要从 UniProt 中选择一个蛋白质作为我们的目标蛋白质。本次实验中，我们将使用人类蛋白质 P63279。下面您可以看到 ESMFold 预测的三维折叠蛋白质结构。
一旦我们有了蛋白质序列，我们就需要定义一些函数来使用我们新训练的 ESM-2 版本，或者原始的 PepMLM。为了简单起见，这里我们只使用原始的 PepMLM，但是您可以轻松地替换成您新微调的 ESM-2 模型。
from transformers import AutoTokenizer, AutoModelForMaskedLM
import torch
import pandas as pd
import numpy as np
from torch.distributions import Categorical
def compute_pseudo_perplexity(model, tokenizer, protein_seq, binder_seq):
sequence = protein_seq + binder_seq
original_input = tokenizer.encode(sequence, return_tensors='pt').to(model.device)
length_of_binder = len(binder_seq)
# Prepare a batch with each row having one masked token from the binder sequence
masked_inputs = original_input.repeat(length_of_binder, 1)
positions_to_mask = torch.arange(-length_of_binder - 1, -1, device=model.device)
masked_inputs[torch.arange(length_of_binder), positions_to_mask] = tokenizer.mask_token_id
# Prepare labels for the masked tokens
labels = torch.full_like(masked_inputs, -100)
labels[torch.arange(length_of_binder), positions_to_mask] = original_input[0, positions_to_mask]
# Get model predictions and calculate loss
with torch.no_grad():
outputs = model(masked_inputs, labels=labels)
loss = outputs.loss
# Loss is already averaged by the model
avg_loss = loss.item()
pseudo_perplexity = np.exp(avg_loss)
return pseudo_perplexity
def generate_peptide_for_single_sequence(protein_seq, peptide_length = 15, top_k = 3, num_binders = 4):
peptide_length = int(peptide_length)
top_k = int(top_k)
num_binders = int(num_binders)
binders_with_ppl = []
for _ in range(num_binders):
# Generate binder
masked_peptide = '<mask>' * peptide_length
input_sequence = protein_seq + masked_peptide
inputs = tokenizer(input_sequence, return_tensors="pt").to(model.device)
with torch.no_grad():
logits = model(**inputs).logits
mask_token_indices = (inputs["input_ids"] == tokenizer.mask_token_id).nonzero(as_tuple=True)[1]
logits_at_masks = logits[0, mask_token_indices]
# Apply top-k sampling
top_k_logits, top_k_indices = logits_at_masks.topk(top_k, dim=-1)
probabilities = torch.nn.functional.softmax(top_k_logits, dim=-1)
predicted_indices = Categorical(probabilities).sample()
predicted_token_ids = top_k_indices.gather(-1, predicted_indices.unsqueeze(-1)).squeeze(-1)
generated_binder = tokenizer.decode(predicted_token_ids, skip_special_tokens=True).replace(' ', '')
# Compute PPL for the generated binder
ppl_value = compute_pseudo_perplexity(model, tokenizer, protein_seq, generated_binder)
# Add the generated binder and its PPL to the results list
binders_with_ppl.append([generated_binder, ppl_value])
return binders_with_ppl
def generate_peptide(input_seqs, peptide_length=15, top_k=3, num_binders=4):
if isinstance(input_seqs, str): # Single sequence
binders = generate_peptide_for_single_sequence(input_seqs, peptide_length, top_k, num_binders)
return pd.DataFrame(binders, columns=['Binder', 'Pseudo Perplexity'])
elif isinstance(input_seqs, list): # List of sequences
results = []
for seq in input_seqs:
binders = generate_peptide_for_single_sequence(seq, peptide_length, top_k, num_binders)
for binder, ppl in binders:
results.append([seq, binder, ppl])
return pd.DataFrame(results, columns=['Input Sequence', 'Binder', 'Pseudo Perplexity'])
model = AutoModelForMaskedLM.from_pretrained("TianlaiChen/PepMLM-650M")
tokenizer = AutoTokenizer.from_pretrained("TianlaiChen/PepMLM-650M")
protein_seq = "MSGIALSRLAQERKAWRKDHPFGFVAVPTKNPDGTMNLMNWECAIPGKKGTPWEGGLFKLRMLFKDDYPSSPPKCKFEPPLFHPNVYPSGTVCLSILEEDKDWRPAITIKQILLGIQELLNEPNIQDPAQAEAYTIYCQNRVEYEKRVRAQAKKFAPS"
results_df = generate_peptide(protein_seq, peptide_length=15, top_k=3, num_binders=5)
print(results_df)
此代码将打印五个可能与我们感兴趣的目标蛋白质结合的肽。具体来说，您应该会看到类似以下内容：
Binder Pseudo Perplexity
0 FQEEPPPLRLAALLL 11.627843
1 EDEDDPEPRYALELE 9.462782
2 FDEDDPLAPRLLEEE 8.092850
3 EQEDPPLPLYALAEE 11.335873
4 FDGEPPLARRLLAKL 10.967774
选择困惑度最低的 FDEDDPLAPRLLEEE，我们可以查询 ESMFold 来预测目标蛋白的结构，该结构通过一个长而灵活的由 20 个甘氨酸组成的连接器与肽结合剂连接。
正如我们所看到的，结构的可信度实际上相当高，但是肽周围的可信度有一些黄色区域，表示可信度较低。我们可以通过使用 EvoProtGrad 对肽区域进行定向进化来提高可信度。
使用 EvoProtGrad 和 ESM-2 对肽结合剂进行计算机指导进化
接下来，我们将（在计算机中）仅对肽结合剂进行定向进化，以尝试提高其对目标蛋白质的结合亲和力。
import evo_prot_grad
from transformers import AutoTokenizer, EsmForMaskedLM
def run_evo_prot_grad_on_paired_sequence(paired_protein_sequence):
# Replace ':' with a string of 20 'G' amino acids
separator = 'G' * 20
sequence_with_separator = paired_protein_sequence.replace(':', separator)
# Determine the start and end indices of the first protein and the separator
separator_start_index = sequence_with_separator.find(separator)
first_protein_end_index = separator_start_index
separator_end_index = separator_start_index + len(separator)
# Format the sequence into FASTA format
fasta_format_sequence = f">Paired_Protein_Sequence\n{sequence_with_separator}"
# Save the sequence to a temporary file
temp_fasta_path = "temp_paired_sequence.fasta"
with open(temp_fasta_path, "w") as file:
file.write(fasta_format_sequence)
# Load the ESM-2 model and tokenizer as the expert
esm2_expert = evo_prot_grad.get_expert(
'esm',
model=EsmForMaskedLM.from_pretrained("facebook/esm2_t33_650M_UR50D"),
tokenizer=AutoTokenizer.from_pretrained("facebook/esm2_t33_650M_UR50D"),
temperature=0.95,
device='cuda' # or 'cpu' if GPU is not available
)
# Initialize Directed Evolution with the preserved first protein and separator region
directed_evolution = evo_prot_grad.DirectedEvolution(
wt_fasta=temp_fasta_path,
output='best',
experts=[esm2_expert],
parallel_chains=1,
n_steps=50,
max_mutations=15,
verbose=True,
preserved_regions=[(0, first_protein_end_index), (separator_start_index, separator_end_index)] # Preserve the first protein and the 'G' amino acids string
)
# Run the evolution process
variants, scores = directed_evolution()
# Process the results and split them into Protein 1 and Protein 2
for variant, score in zip(variants, scores):
# Remove spaces from the sequence
evolved_sequence_no_spaces = variant.replace(" ", "")
# Split the sequence at the separator
protein_1, protein_2 = evolved_sequence_no_spaces.split(separator)
print(f"Protein: {protein_1}, Evolved Peptide: {protein_2}, Score: {score}")
# Example usage
paired_protein_sequence = "MSGIALSRLAQERKAWRKDHPFGFVAVPTKNPDGTMNLMNWECAIPGKKGTPWEGGLFKLRMLFKDDYPSSPPKCKFEPPLFHPNVYPSGTVCLSILEEDKDWRPAITIKQILLGIQELLNEPNIQDPAQAEAYTIYCQNRVEYEKRVRAQAKKFAPS:FDEDDPLAPRLLEEE" # Replace with your paired protein sequences
run_evo_prot_grad_on_paired_sequence(paired_protein_sequence)
这将打印使用ESM-2作为专家模型的基于梯度的MCMC方法的50次迭代。
>Wildtype sequence: M S G I A L S R L A Q E R K A W R K D H P F G F V A V P T K N P D G T M N L M N W E C A I P G K K G T P W E G G L F K L R M L F K D D Y P S S P P K C K F E P P L F H P N V Y P S G T V C L S I L E E D K D W R P A I T I K Q I L L G I Q E L L N E P N I Q D P A Q A E A Y T I Y C Q N R V E Y E K R V R A Q A K K F A P S G G G G G G G G G G G G G G G G G G G G F D E D D P L A P R L L E E E
step 0 acceptance rate: 0.1181
>chain 0, Product of Experts score: 0.0000
M S G I A L S R L A Q E R K A W R K D H P F G F V A V P T K N P D G T M N L M N W E C A I P G K K G T P W E G G L F K L R M L F K D D Y P S S P P K C K F E P P L F H P N V Y P S G T V C L S I L E E D K D W R P A I T I K Q I L L G I Q E L L N E P N I Q D P A Q A E A Y T I Y C Q N R V E Y E K R V R A Q A K K F A P S G G G G G G G G G G G G G G G G G G G G F D E D D P L A P R L L E E E
step 1 acceptance rate: 0.3120
>chain 0, Product of Experts score: 0.0000
M S G I A L S R L A Q E R K A W R K D H P F G F V A V P T K N P D G T M N L M N W E C A I P G K K G T P W E G G L F K L R M L F K D D Y P S S P P K C K F E P P L F H P N V Y P S G T V C L S I L E E D K D W R P A I T I K Q I L L G I Q E L L N E P N I Q D P A Q A E A Y T I Y C Q N R V E Y E K R V R A Q A K K F A P S G G G G G G G G G G G G G G G G G G G G F D E D D P L A V R L L E E E
step 2 acceptance rate: 0.0463
>chain 0, Product of Experts score: 1.8138
M S G I A L S R L A Q E R K A W R K D H P F G F V A V P T K N P D G T M N L M N W E C A I P G K K G T P W E G G L F K L R M L F K D D Y P S S P P K C K F E P P L F H P N V Y P S G T V C L S I L E E D K D W R P A I T I K Q I L L G I Q E L L N E P N I Q D P A Q A E A Y T I Y C Q N R V E Y E K R V R A Q A K K F A P S G G G G G G G G G G G G G G G G G G G G F D E D D P L A V R L F E E E
step 3 acceptance rate: 0.1111
.
.
.
>chain 0, Product of Experts score: -2.7335
M S G I A L S R L A Q E R K A W R K D H P F G F V A V P T K N P D G T M N L M N W E C A I P G K K G T P W E G G L F K L R M L F K D D Y P S S P P K C K F E P P L F H P N V Y P S G T V C L S I L E E D K D W R P A I T I K Q I L L G I Q E L L N E P N I Q D P A Q A E A Y T I Y C Q N R V E Y E K R V R A Q A K K F A P S G G G G G G G G G G G G G G G G G G G G F V G I D F S P R V G E D K V
step 48 acceptance rate: 0.0667
>chain 0, Product of Experts score: -2.7335
M S G I A L S R L A Q E R K A W R K D H P F G F V A V P T K N P D G T M N L M N W E C A I P G K K G T P W E G G L F K L R M L F K D D Y P S S P P K C K F E P P L F H P N V Y P S G T V C L S I L E E D K D W R P A I T I K Q I L L G I Q E L L N E P N I Q D P A Q A E A Y T I Y C Q N R V E Y E K R V R A Q A K K F A P S G G G G G G G G G G G G G G G G G G G G F V G I D F S M R V G E D K V
step 49 acceptance rate: 1.9997
>chain 0, Product of Experts score: -4.7062
M S G I A L S R L A Q E R K A W R K D H P F G F V A V P T K N P D G T M N L M N W E C A I P G K K G T P W E G G L F K L R M L F K D D Y P S S P P K C K F E P P L F H P N V Y P S G T V C L S I L E E D K D W R P A I T I K Q I L L G I Q E L L N E P N I Q D P A Q A E A Y T I Y C Q N R V E Y E K R V R A Q A K K F A P S G G G G G G G G G G G G G G G G G G G G F V H I D F S M R V G C D K V
Protein: MSGIALSRLAQERKAWRKDHPFGFVAVPTKNPDGTMNLMNWECAIPGKKGTPWEGGLFKLRMLFKDDYPSSPPKCKFEPPLFHPNVYPSGTVCLSILEEDKDWRPAITIKQILLGIQELLNEPNIQDPAQAEAYTIYCQNRVEYEKRVRAQAKKFAPS, Evolved Peptide:
FSPSDVNKLFGKDED, Score: 2.326873779296875
请注意，此过程实际上**增加了**困惑度。
model = AutoModelForMaskedLM.from_pretrained("TianlaiChen/PepMLM-650M")
tokenizer = AutoTokenizer.from_pretrained("TianlaiChen/PepMLM-650M")
protein_seq = "MSGIALSRLAQERKAWRKDHPFGFVAVPTKNPDGTMNLMNWECAIPGKKGTPWEGGLFKLRMLFKDDYPSSPPKCKFEPPLFHPNVYPSGTVCLSILEEDKDWRPAITIKQILLGIQELLNEPNIQDPAQAEAYTIYCQNRVEYEKRVRAQAKKFAPS"
binder_seq = "FSPSDVNKLFGKDED"
compute_pseudo_perplexity(model, tokenizer, protein_seq, binder_seq)
21.214165484401022
然而，我们看到 ESMFold 对预测结构的可信度实际上增加了。这表明，仅仅根据困惑度来判断结合剂的质量可能是一个不够充分的指标。
放大查看肽结合剂与目标蛋白的相互作用，我们可以看到目标蛋白与肽结合剂之间的原子距离实际上非常接近。
我们还注意到，在计算困惑度时应考虑一些细微差别。特别是，在自然语言处理领域已经表明，由内在维度衡量的几何压缩实际上可以预测困惑度（例如，参见 使用 ESM-2 估计蛋白质序列嵌入的内在维度 和 弥合语言模型中的信息论和几何压缩）。这为我们提供了关于在训练或微调过程中模型适应难度的描述，并且可以为我们提供衡量蛋白质或肽“自然”程度的指标，但它不一定能为我们提供衡量结合剂在结合亲和力方面的优劣的指标。为此，我们可能需要另一种专门用于预测结合亲和力的方法。
结论
对 ESM-2 及其在生成和优化目标蛋白肽结合剂方面的应用的探索，揭示了人工智能驱动方法在蛋白质工程领域的巨大潜力。通过训练 PepMLM、生成肽结合剂以及使用 EvoProtGrad 进行计算机定向进化的过程，我们展示了一种增强肽-蛋白质相互作用的综合方法。ESM-2 的使用不仅促进了潜在肽结合剂的识别，还允许对其进行后续优化，这表明蛋白质科学在预测和设计能力方面取得了重大进展。这种方法有望应用于广泛的领域，包括药物发现、治疗性蛋白质设计和理解复杂的生物相互作用。本研究中人工智能与计算生物学的结合，代表着分子生物学和生物工程领域更高效、更有针对性方法的重大进步。有关 PepMLM 的更多信息，您可以阅读研究论文。有关 EvoProtGrad 的更多信息，请查阅研究论文和Github。