基于对比学习的酶功能预测：从算法原理到生产实践

共计 4227 个字符，预计需要花费 11 分钟才能阅读完成。

背景痛点

酶功能预测一直是生物信息学中的重要挑战。传统的 BLAST 和 Pfam 方法虽然被广泛使用，但它们存在明显的局限性：

• 依赖序列比对，对新酶家族的泛化能力较差
• 对远缘同源酶的识别准确率低
• 无法有效捕捉序列 - 功能关系的深层次特征

这些限制促使我们寻找更先进的机器学习方法来解决这个问题。

技术对比

核心实现

1. 使用 ESM- 2 预训练模型

ESM- 2 是目前最先进的蛋白质语言模型之一，能够为任意蛋白序列生成高质量的 embedding 表示。

import torch
from esm import pretrained

# 加载预训练的 ESM- 2 模型
model, alphabet = pretrained.load_model_and_alphabet('esm2_t33_650M_UR50D')
batch_converter = alphabet.get_batch_converter()

# 将蛋白序列转换为 embedding
data = [("seq1", "MKTVRQERL...")]  # 输入序列
batch_labels, batch_strs, batch_tokens = batch_converter(data)
with torch.no_grad():
    results = model(batch_tokens, repr_layers=[33])
# 获取最后一层的 embedding
embedding = results["representations"][33]

2. 构建正负样本对

正样本对选择同一 EC 编号下的不同酶，负样本对选择不同 EC 编号的酶。我们采用 SimCLR 框架进行对比学习。

class ContrastiveDataset(torch.utils.data.Dataset):
    def __init__(self, fasta_file, ec_mapping):
        """
        初始化对比学习数据集
        :param fasta_file: FASTA 格式的蛋白序列文件
        :param ec_mapping: 序列到 EC 编号的映射
        """
        self.sequences = self._load_fasta(fasta_file)
        self.ec_dict = ec_mapping
        # 构建 EC 编号到序列列表的倒排索引
        self.ec_to_seqs = self._build_ec_index()

    def _load_fasta(self, filepath):
        # 实现 FASTA 文件加载
        pass

    def _build_ec_index(self):
        # 构建 EC 编号索引
        pass

    def __getitem__(self, idx):
        # 返回一个正样本对和若干个负样本
        anchor_seq = self.sequences[idx]
        anchor_ec = self.ec_dict[anchor_seq.id]

        # 正样本：同一 EC 下的不同序列
        pos_seqs = [s for s in self.ec_to_seqs[anchor_ec] if s.id != anchor_seq.id]
        pos_seq = random.choice(pos_seqs)

        # 负样本：不同 EC 的序列
        all_ecs = list(self.ec_to_seqs.keys())
        neg_ecs = [ec for ec in all_ecs if ec != anchor_ec]
        neg_ec = random.choice(neg_ecs)
        neg_seq = random.choice(self.ec_to_seqs[neg_ec])

        return anchor_seq, pos_seq, neg_seq

3. 损失函数与训练

我们使用 NT-Xent(Normalized Temperature-scaled Cross Entropy) 作为损失函数，并采用梯度累积来支持更大的 batch size。

class ContrastiveLoss(nn.Module):
    def __init__(self, temperature=0.07):
        """
        NT-Xent 损失函数
        :param temperature: ** 温度系数 **，控制样本分布的尖锐程度
        """
        super().__init__()
        self.temperature = temperature
        self.criterion = nn.CrossEntropyLoss()

    def forward(self, features):
        """
        :param features: 模型输出的特征，形状为 (2N, D)
                        其中每两个连续样本构成一个正对
        """
        batch_size = features.shape[0]

        # 计算相似度矩阵
        similarity = nn.functional.cosine_similarity(features.unsqueeze(1), features.unsqueeze(0), dim=2
        ) / self.temperature

        # 构建标签：每个样本的正样本是它的配对样本
        labels = torch.arange(batch_size, device=features.device)
        labels = (labels + 1 - labels % 2 * 2)  # 1->0, 0->1, 3->2, 2->3,...

        # 计算损失
        loss = self.criterion(similarity, labels)
        return loss

# 训练循环示例
model = Model()  # 自定义模型
optimizer = torch.optim.Adam(model.parameters(), lr=3e-4)
criterion = ContrastiveLoss(temperature=0.05)  # ** 重要超参数 **

# 梯度累积步骤
accum_steps = 4  # ** 梯度累积次数 **

for epoch in range(100):
    for i, (anchor, pos, neg) in enumerate(train_loader):
        # 前向传播
        anchor_emb = model(anchor)
        pos_emb = model(pos)
        neg_emb = model(neg)

        # 拼接所有样本特征
        features = torch.cat([anchor_emb.unsqueeze(1), 
                             pos_emb.unsqueeze(1), 
                             neg_emb.unsqueeze(1)], dim=1)
        features = features.view(-1, features.size(-1))

        # 计算损失
        loss = criterion(features)

        # 反向传播 (梯度累积)
        loss = loss / accum_steps
        loss.backward()

        if (i + 1) % accum_steps == 0:
            optimizer.step()
            optimizer.zero_grad()

生产考量

1. 处理 EC 编号层级结构

EC 编号是层级化的 (如 EC 1.2.3.4)，我们可以利用这种结构信息：

def ec_hierarchy_loss(predictions, targets):
    """
    考虑 EC 编号层级结构的损失函数
    第一级错误比第四级错误惩罚更重
    """
    # 拆分 EC 编号的各级
    pred_levels = predictions.split('.', maxsplit=3)
    target_levels = targets.split('.', maxsplit=3)

    loss = 0
    weights = [0.4, 0.3, 0.2, 0.1]  # 层级权重
    for i, (p, t) in enumerate(zip(pred_levels, target_levels)):
        if p != t:
            loss += weights[i]
            break
    return loss

2. 分布式训练策略

在多 GPU 训练时，需要同步各设备的 embedding 计算：

import torch.distributed as dist

def gather_embeddings(embeddings):
    """在所有进程中收集 embedding"""
    gathered_embeddings = [torch.zeros_like(embeddings) 
                          for _ in range(dist.get_world_size())]
    dist.all_gather(gathered_embeddings, embeddings)
    return torch.cat(gathered_embeddings)

3. 模型解释性

使用 SHAP 值分析模型预测的依据：

import shap

# 创建解释器
explainer = shap.DeepExplainer(model, background_data)

# 计算单个序列的 SHAP 值
shap_values = explainer.shap_values(test_sequence)

# 可视化
shap.initjs()
shap.force_plot(explainer.expected_value[0], 
                shap_values[0], 
                feature_names=amino_acids)

避坑指南

1. 避免序列相似性泄漏 ：
2. 使用严格的聚类划分：将序列相似性 >30% 的序列划分到同一 fold

3. 确保训练集和测试集没有高度相似的序列

4. 显存优化 ：

   # 启用混合精度训练
   scaler = torch.cuda.amp.GradScaler()
   
   with torch.cuda.amp.autocast():
       embeddings = model(sequences)
       loss = criterion(embeddings)
   
   scaler.scale(loss).backward()
   scaler.step(optimizer)
   scaler.update()

5. 负样本平衡 ：

6. 对罕见的 EC 类别进行过采样
7. 使用动态负采样策略，根据训练过程中的表现调整采样频率

结论与开放问题

本文介绍的对比学习方法显著提升了酶功能预测的准确率，特别是在处理新酶家族时表现优异。但仍有一些开放问题值得探索：

1. 如何有效整合 AlphaFold 预测的蛋白结构信息？
2. 能否设计更智能的负样本采样策略，进一步提升模型性能？
3. 对于多功能的酶 (具有多个 EC 编号)，如何改进模型架构？

希望这篇实践指南能为生物信息学和机器学习交叉领域的研究者提供有价值的参考。完整的实现代码已开源在 GitHub 上，欢迎交流讨论。