使用持久同源性聚类蛋白质复合物并微调ESM-2用于PPI网络预测

蛋白质复合物的聚类目前仍是一个开放性问题。本文将介绍一种基于蛋白质语言模型嵌入和持久同源性的蛋白质序列和蛋白质复合物聚类新方法。然后，我们将展示如何使用这种新的聚类方法基于训练/测试拆分来微调ESM-2以预测蛋白质-蛋白质相互作用。

蛋白质复合物

蛋白质是细胞的“主力军”，它们通常不单独行动。许多重要的生物功能是由蛋白质复合物完成的，这些复合物是由两个或多个蛋白质分子结合形成的结构。这些复合物的形成是一个高度受调节且特异的过程，对于许多细胞活动至关重要，包括信号通路、代谢反应、DNA复制等等。理解蛋白质复合物对于从分子水平上掌握细胞机制至关重要。这里我们将重点关注由两种蛋白质组成的蛋白质复合物，但这种方法可以扩展到由多种蛋白质组成的复合物。

序列相似性

测量两个蛋白质序列之间的相似性以及蛋白质同源建模可以通过多种方式完成。其中最古老的方法之一是基于编辑距离，但也有较新的方法利用蛋白质语言模型的嵌入来确定两个蛋白质的相似程度。例如，在论文蛋白质语言模型实现高效同源检测中，作者设计了一种使用较旧的ESM-1b模型嵌入进行相似性搜索的方法。我们从这些方法中汲取灵感，并在我们的方法中使用较新的蛋白质语言模型ESM-2。我们还使用拓扑数据分析技术来获得与蛋白质和蛋白质复合物相关的嵌入的拓扑摘要。虽然这种方法可以用于单个蛋白质，但我们更关注它如何应用于蛋白质复合物，因为当试图确定两个蛋白质复合物的相似性时，标准的序列相似性度量会变得不足或过于复杂。事实上，大多数蛋白质-蛋白质复合物同源建模方法都是基于结构的方法。我们将尝试通过一种称为持久同源性的方法来弥补这一点，该方法应用于使用蛋白质语言模型ESM-2的蛋白质复合物嵌入。

蛋白质复合物聚类

蛋白质复合物的聚类是一个难题，我们将使用持久同源性以新颖的方式来解决。我们首先连接从UniProt数据库中获得的人类相互作用蛋白质对。接下来，我们计算pLM ESM-2关联到连接对的嵌入。一旦我们有了嵌入，我们就会使用持久同源性计算一个叫做持久图的东西。每个蛋白质复合物都给出了这样一个图，并且这些图之间的成对距离是使用持久图上的Wasserstein距离度量计算的。接下来，我们使用获得的距离矩阵计算一个二级持久图，该图总结了与每个蛋白质-蛋白质复合物相关的持久图之间的Wasserstein距离。然后，我们根据这个二级持久图运行DBSCAN，选择一个ε阈值，使得（二级）零维持久图中80%的点落在ε以下。这将返回蛋白质-蛋白质复合物的簇，然后我们可以使用这些簇来创建训练/测试集，用于训练模型来预测蛋白质-蛋白质相互作用或生成靶蛋白的结合剂。

聚类脚本

下面，我们提供了基于持久同源性聚类蛋白质-蛋白质复合物的脚本。要使用它，您可以从此处HuggingFace下载相互作用蛋白质数据集。

import pandas as pd
import numpy as np
from transformers import EsmModel, AutoTokenizer
import torch
from scipy.spatial.distance import pdist, squareform
from gudhi import RipsComplex
from gudhi.hera import wasserstein_distance
from sklearn.cluster import DBSCAN
from sklearn.metrics import silhouette_score
from tqdm import tqdm

# Define a helper function for hidden states
def get_hidden_states(sequence, tokenizer, model, layer):
    model.config.output_hidden_states = True
    encoded_input = tokenizer([sequence], return_tensors='pt', padding=True, truncation=True, max_length=1024)
    with torch.no_grad():
        model_output = model(**encoded_input)
    hidden_states = model_output.hidden_states
    specific_hidden_states = hidden_states[layer][0]
    return specific_hidden_states.numpy()

# Define a helper function for Euclidean distance matrix
def compute_euclidean_distance_matrix(hidden_states):
    euclidean_distances = pdist(hidden_states, metric='euclidean')
    euclidean_distance_matrix = squareform(euclidean_distances)
    return euclidean_distance_matrix

# Define a helper function for persistent homology
def compute_persistent_homology(distance_matrix, max_dimension=0):
    max_edge_length = np.max(distance_matrix)
    rips_complex = RipsComplex(distance_matrix=distance_matrix, max_edge_length=max_edge_length)
    st = rips_complex.create_simplex_tree(max_dimension=max_dimension)
    st.persistence()
    return st, st.persistence()

# Define a helper function for Wasserstein distances
def compute_wasserstein_distances(persistence_diagrams, dimension):
    n_diagrams = len(persistence_diagrams)
    distances = np.zeros((n_diagrams, n_diagrams))
    for i in tqdm(range(n_diagrams), desc="Computing Wasserstein Distances"):
        for j in range(i+1, n_diagrams):
            X = np.array([p[1] for p in persistence_diagrams[i] if p[0] == dimension])
            Y = np.array([p[1] for p in persistence_diagrams[j] if p[0] == dimension])
            distance = wasserstein_distance(X, Y)
            distances[i][j] = distance
            distances[j][i] = distance
    return distances

# Load the tokenizer and model
tokenizer = AutoTokenizer.from_pretrained("facebook/esm2_t33_650M_UR50D")
model = EsmModel.from_pretrained("facebook/esm2_t33_650M_UR50D")

# Define layer to be used
layer = model.config.num_hidden_layers - 1

# Load the TSV file
file_path = 'pepmlm/scripts/data/filtered_protein_interaction_pairs.tsv'
protein_pairs_df = pd.read_csv(file_path, sep='\t')

# Only process the first 1000 proteins
protein_pairs_df = protein_pairs_df.head(1000)

# Extract concatenated sequences
concatenated_sequences = protein_pairs_df['Protein1'] + protein_pairs_df['Protein2']

# Initialize list to store persistent diagrams
persistent_diagrams = []

# Loop over concatenated sequences to compute their persistent diagrams
for sequence in tqdm(concatenated_sequences, desc="Computing Persistence Diagrams"):
    hidden_states_matrix = get_hidden_states(sequence, tokenizer, model, layer)
    distance_matrix = compute_euclidean_distance_matrix(hidden_states_matrix)
    _, persistence_diagram = compute_persistent_homology(distance_matrix)
    persistent_diagrams.append(persistence_diagram)

# Compute the Wasserstein distances
wasserstein_distances = compute_wasserstein_distances(persistent_diagrams, 0)

# Compute the second-level persistent homology
with tqdm(total=1, desc="Computing Second-Level Persistent Homology") as pbar:
    st_2, persistence_2 = compute_persistent_homology(wasserstein_distances)
    pbar.update(1)

# Function to calculate the epsilon for DBSCAN
def calculate_epsilon(persistence_diagrams, threshold_percentage):
    lifetimes = [p[1][1] - p[1][0] for p in persistence_diagrams if p[0] == 0]
    lifetimes.sort()
    threshold_index = int(threshold_percentage * len(lifetimes))
    return lifetimes[threshold_index]

# Calculate epsilon
threshold_percentage = 0.8  # 80%
epsilon = calculate_epsilon(persistence_2, threshold_percentage)

# Perform DBSCAN clustering
with tqdm(total=1, desc="Performing DBSCAN Clustering") as pbar:
    dbscan = DBSCAN(metric="precomputed", eps=epsilon, min_samples=1)
    dbscan.fit(wasserstein_distances)
    labels = dbscan.labels_
    pbar.update(1)

# Add the cluster labels to the DataFrame
protein_pairs_df['Cluster'] = labels

# Save the DataFrame with cluster information
output_file_path = 'clustered_protein_interaction_pairs.tsv'
protein_pairs_df.to_csv(output_file_path, sep='\t', index=False)

print(f"Clustered data saved to: {output_file_path}")

在此脚本中，我们对蛋白质-蛋白质复合物的持久同源性进行了以下解释和使用：

1. 持久图和生命周期： 在脚本中，持久图是使用拓扑数据分析（TDA）技术为每个连接的蛋白质序列生成的。这些图捕获了拓扑特征（如簇或循环）在各种尺度上的“诞生”和“消亡”。特征的“生命周期”是其消亡值与诞生值之差。

2. 二级持久图： 计算所有持久图对之间的Wasserstein距离后，生成一个二级持久图。此图表示Wasserstein距离空间中特征（在本例中为簇）的持久性。

3. 计算生命周期： 对于二级持久图零维部分的每个点，都计算其生命周期。这些生命周期表示Wasserstein距离空间中簇的持久性。

4. 确定80%阈值： 对生命周期进行排序，并确定第80个百分位数。这意味着零维持久图中80%的点的生命周期小于或等于此值。

5. DBSCAN中的ε设置： 将80%的阈值用作DBSCAN中的epsilon参数。在DBSCAN中，epsilon确定两个样本之间被视为彼此邻域的最大距离。通过将epsilon设置为此80%阈值，聚类算法被调整为捕获拓扑分析所揭示的大多数自然聚类结构。

6. 解释： 使用80%的阈值设置epsilon意味着DBSCAN聚类对数据中最持久和最重要的特征敏感，这些特征由二级持久图捕获。这种方法旨在识别在数据拓扑结构背景下稳健且重要的簇。

由于以这种方式聚类蛋白质-蛋白质复合物的计算成本过高，我们将仅聚类前1000个蛋白质序列。寻找优化这种聚类方法的方法将允许聚类更多的蛋白质-蛋白质复合物。作为替代方案，我们可以遵循蛋白质语言模型实现高效同源检测中提及的方法。这可能也有助于我们改进此博客文章中提到的vcMSA的方法。我们还可以使用GPU计算持久同源性来加快速度。

接下来，我们根据这些簇创建一个训练/测试拆分，这将发挥类似于基于序列相似性聚类序列或基于UniProt数据库中蛋白质家族创建训练/测试拆分的作用。然而，这种形式的聚类是基于对象，这些对象是从嵌入向量中获得的，称为过滤单纯复形，它们是蛋白质-蛋白质复合物的一种几何指纹。它们是编码在嵌入向量（ESM-2的隐藏状态）中的语义信息的拓扑摘要。

在相互作用蛋白质对上微调ESM-2

接下来，我们可以在UniProt的相互作用蛋白质对上微调ESM-2，以提高其建模蛋白质复合物、预测蛋白质-蛋白质相互作用或生成结合剂的能力。我们通过运行以下脚本，对上一节中获得的人类聚类蛋白质进行一个 epoch 的微调。此脚本将微调ESM-2以基于MLM损失预测PPI网络，如这篇博客文章中提到的。我们将在后续文章中介绍如何微调ESM-2以生成结合剂，这将类似于这篇博客文章中提到的微调ESM-2以获得PepMLM的方法。

from transformers import Trainer, TrainingArguments, AutoTokenizer, EsmForMaskedLM, TrainerCallback, get_scheduler
from torch.utils.data import Dataset
import pandas as pd
import torch
from torch.optim import AdamW
from sklearn.model_selection import train_test_split
import random

class ProteinDataset(Dataset):
    def __init__(self, proteins, peptides, tokenizer):
        self.tokenizer = tokenizer
        self.proteins = proteins
        self.peptides = peptides

    def __len__(self):
        return len(self.proteins)

    def mask_sequence(self, sequence, mask_percentage):
        mask_indices = random.sample(range(len(sequence)), int(len(sequence) * mask_percentage))
        return ''.join([self.tokenizer.mask_token if i in mask_indices else char for i, char in enumerate(sequence)])

    def __getitem__(self, idx):
        protein_seq = self.proteins[idx]
        peptide_seq = self.peptides[idx]

        masked_protein = self.mask_sequence(protein_seq, 0.55)
        masked_peptide = self.mask_sequence(peptide_seq, 0.55)
        complex_seq = masked_protein + masked_peptide

        # Tokenize and pad the complex sequence
        complex_input = self.tokenizer(
            complex_seq, 
            return_tensors="pt", 
            padding="max_length", 
            max_length=1024, 
            truncation=True, 
            add_special_tokens=False
        )

        input_ids = complex_input["input_ids"].squeeze()
        attention_mask = complex_input["attention_mask"].squeeze()

        # Create labels
        label_seq = protein_seq + peptide_seq
        labels = self.tokenizer(
            label_seq, 
            return_tensors="pt", 
            padding="max_length", 
            max_length=1024, 
            truncation=True, 
            add_special_tokens=False
        )["input_ids"].squeeze()

        # Set non-masked positions in the labels tensor to -100
        labels = torch.where(input_ids == self.tokenizer.mask_token_id, labels, -100)

        return {"input_ids": input_ids, "attention_mask": attention_mask, "labels": labels}


# Loading the dataset
file_path = "clustered_protein_interaction_pairs.tsv"
data = pd.read_csv(file_path, delimiter='\t')

# Splitting the data based on clusters
cluster_sizes = data['Cluster'].value_counts(normalize=True)
test_clusters = []
test_size = 0

for cluster, size in cluster_sizes.items():
    test_clusters.append(cluster)
    test_size += size
    if test_size >= 0.20:
        break

test_data = data[data['Cluster'].isin(test_clusters)]
train_data = data[~data['Cluster'].isin(test_clusters)]

proteins_train = train_data["Protein1"].tolist()
peptides_train = train_data["Protein2"].tolist()
proteins_test = test_data["Protein1"].tolist()
peptides_test = test_data["Protein2"].tolist()

# Load tokenizer and model
model_name = "esm2_t30_150M_UR50D"
tokenizer = AutoTokenizer.from_pretrained("facebook/" + model_name)
model = EsmForMaskedLM.from_pretrained("facebook/" + model_name)

# Training arguments
training_args = TrainingArguments(
    output_dir='./interact_output/',
    num_train_epochs=3,
    per_device_train_batch_size=1,
    per_device_eval_batch_size=4,
    warmup_steps=50,
    logging_dir='./logs',
    logging_steps=10,
    evaluation_strategy="epoch",
    load_best_model_at_end=True,
    save_strategy='epoch',
    metric_for_best_model='eval_loss',
    save_total_limit=3,
    gradient_accumulation_steps=2,
    lr_scheduler_type='cosine'  # Set learning rate scheduler to cosine
)

# Optimizer
optimizer = AdamW(model.parameters(), lr=0.0007984276816171436)

# Scheduler
scheduler = get_scheduler(
    name="cosine",
    optimizer=optimizer,
    num_warmup_steps=training_args.warmup_steps,
    num_training_steps=training_args.max_steps
)

# Instantiate the ProteinDataset for training and testing
train_dataset = ProteinDataset(proteins_train, peptides_train, tokenizer)
test_dataset = ProteinDataset(proteins_test, peptides_test, tokenizer)

# Trainer
trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=train_dataset,
    eval_dataset=test_dataset,
    optimizers=(optimizer, scheduler),
)

# Start training
trainer.train()

结论

现在您已经拥有经过微调的ESM-2模型用于预测蛋白质-蛋白质相互作用，您可以按照这篇博客文章来预测PPI网络。您还可以尝试调整训练脚本以匹配PepMLM的训练脚本，以便结合剂被完全遮罩，而靶序列保持未遮罩，从而根据这篇文章中解释的方法，微调ESM-2以生成靶蛋白的结合剂，或者您可以等待下一篇博客文章，它很可能会涵盖这一点！虽然这种方法从数学角度来看是新颖而有趣的，但使用CPU进行持久同源性聚类似乎在计算上过于昂贵。也许作为下一个项目，您可以尝试使用这个在GPU上实现持久同源性计算，以加快计算速度！