基于似然函数的蛋白质语言模型微调用于少样本适应性预测和设计

这篇文章讲述如何有效地微调蛋白质语言模型 (Protein Language Models, PLMs)，以提升其在少样本 (few-shot) 条件下的蛋白质适应性 (fitness) 预测和序列设计能力。

1. 背景知识

自然界中的蛋白质序列是长期进化和自然选择的结果。蛋白质语言模型 (PLMs) 通过学习大量自然蛋白质序列的数据，能够隐式地捕捉到蛋白质序列中蕴含的分布约束，这些约束与蛋白质的功能和结构息息相关。因此，PLMs 可以作为有效的零样本 (zero-shot) 预测器，直接预测氨基酸突变对蛋白质适应性的影响。

近年来，研究人员提出了多种利用 PLMs 蕴含的分布知识来提升有监督的蛋白质适应性预测和序列设计任务的方法。然而，由于缺乏对不同预测策略和不同类型 PLM 模型的全面比较，目前尚不清楚哪种方法性能最佳。

本文着重解决这个问题，并在以下几个方面做出了贡献：

• 扩展了基于排序的损失函数，用于开发基于似然 (likelihood) 的评分函数，从而微调掩码语言模型 (masked PLMs) 和家族式语言模型 (family-based PLMs)。
• 直接对比了最先进的 (SOTA) 基于冻结嵌入 (frozen embeddings) 的方法 以及基于回归的微调方法，证明了排序损失微调方案的有效性，尤其是在数据量较少的情况下。
• 开发了与排序损失微调方案兼容的集成 (ensemble) 策略，并证明了其在多轮贝叶斯优化 (Bayesian optimisation) 设置中的有效性。

总而言之，本文旨在探索如何更好地利用 PLMs 的零样本预测能力，结合少量有标签数据，来提升蛋白质适应性预测和序列设计的性能，特别是在数据稀缺的实际场景中。

2. 文章概述

本文研究了如何通过微调蛋白质语言模型 (PLMs) 来提高少样本 (few-shot) 蛋白质适应性预测和序列设计的效果。作者扩展了先前提出的基于排序的损失函数，为掩码 PLM (如 ESM-1v, ESM-2) 和家族式自回归 PLM (如 POET) 开发了基于似然的评分函数。

文章首先对比了基于排序损失的微调和基于均方误差 (MSE) 的回归微调，以及基于冻结嵌入的 SOTA 方法 (ProteinNPT)。实验结果表明，在低数据量情况下，基于排序损失的微调方法优于其他方法，尤其对于家族式 PLM 模型 POET 提升显著。

此外，作者还提出了集成策略，利用 PLMs 学习到的突变分布对序列上下文的强依赖性，进一步提升预测性能，并应用于指导蛋白质适应性景观 (fitness landscapes) 的高效优化。

实验结果表明，本文提出的基于似然函数和排序损失的微调方法，以及集成策略，在少样本蛋白质适应性预测和序列设计任务中，相较于现有方法，能够取得更优的性能。

3. 方法详解

3.1. 数据集

本文使用了 ProteinGym 数据集，这是一个包含多个蛋白质突变景观 (mutational landscapes) 的数据集，每个景观包含一组蛋白质序列以及实验测定的适应性数值。

• 单突变景观 (Single-mutant landscapes): 使用了由 Notin 等人 (2023b) 选出的 8 个代表性的单突变景观验证集。
• 多突变景观 (Multiple-mutant landscapes): 使用了 5 个包含多重突变的景观，这些景观来自 ProteinGym 中多样且非冗余的景观子集。

对于每个景观，作者在 低数据量 (low-data) 设置下进行实验，训练集大小 n 分别为 32, 128 或 512 个序列，评估集大小为单突变景观 2000 个序列，多突变景观 5000 个序列。此外，使用额外的 128 个随机采样的序列作为验证集，用于提前停止 (early stopping)。

3.2. 算法

本文主要使用了以下几种蛋白质语言模型 (PLMs):

• 掩码语言模型 (Masked PLMs):
  • ESM-1v (650M) (Meier et al., 2021)
  • ESM-2 (650M) (Lin et al., 2022)
• 家族式自回归模型 (Family-based Autoregressive Model):
  • POET (Truong Jr & Bepler, 2023)

并对比了以下基线方法:

• ProteinNPT (MSAT, ESM-1v) (Notin et al., 2023b): 一种基于冻结 PLM 嵌入的 SOTA 方法。
• Augmented Density Models (OHE aug., Emb. aug.): 增强密度模型，结合 PLM 的零样本预测和序列表示 (one-hot 编码或 PLM 嵌入)。

微调策略 (Fine-tuning Strategies):

本文的核心方法是使用排序损失函数 (Ranking-based loss functions) 来微调 PLMs 的似然函数，并与以下方法进行对比：

• 基于均方误差的回归微调 (MSE loss): 直接回归预测适应性数值。
• 线性回归头 (Linear Regression Head): 在 PLM 嵌入之上添加线性回归层进行微调。

评分函数 (Scoring Functions):

• 自回归 PLM (Autoregressive PLMs, e.g., POET): 使用条件序列对数似然 (conditional sequence log-likelihood) 作为评分函数:

  $s_\theta(x, M) = \sum_{i=1}^{L} \log p_\theta(x_i | x_{<i}, M)$

  其中, $x$ 是蛋白质序列, $M$ 是多序列比对 (MSA), $L$ 是序列长度, $p_\theta$ 是模型预测的条件概率。

• 掩码语言模型 (Masked PLMs, e.g., ESM-1v, ESM-2): 使用 wild-type marginals (wt-marginals) 评分函数:

  $s_\theta(x, x^{wt}) = \sum_{i \in I^{mut}} [\log p_\theta(x_i | x^{wt}) - \log p_\theta(x^{wt}_i | x^{wt})]$

  其中, $x$ 是突变序列, $x^{wt}$ 是野生型序列, $I^{mut}$ 是突变位点的集合, $p_\theta$ 是掩码语言模型预测的概率。

损失函数 (Loss Functions):

• 排序损失 (Ranking Loss): 使用 Bradley-Terry 模型参数化的 pairwise ranking loss:

  $L = \sum_{i=1}^{B} \sum_{j=1}^{B} - \mathbb{I}(y(x_i) > y(x_j)) \log \sigma(s_\theta(x_i, c) - s_\theta(x_j, c))$

  其中, $x_i, x_j$ 是一个 batch 中的序列, $y(x_i), y(x_j)$ 是对应的适应性数值, $s_\theta(x, c)$ 是评分函数, $\sigma$ 是 sigmoid 函数, $\mathbb{I}$ 是指示函数。

• 均方误差损失 (MSE Loss): 标准的均方误差损失函数，用于回归任务。

3.3. 训练和评估

• 训练 (Training): 使用 Adam 优化器 (Kingma & Ba, 2015) 微调模型，有效 batch size 为 32。针对不同的微调策略 (回归 vs 排序)，分别在单突变景观上扫过学习率 1e-4, 3e-5, 1e-5 后选择最佳学习率。模型训练 50 个 epochs，并使用验证集上的 Spearman 相关系数进行早停。
• 评估 (Evaluation): 使用测试集上的 Spearman 相关系数 (Spearman correlation) 评估适应性预测性能。对于序列设计任务，使用 Recall 指标 (Top 100 Recall, Top 30% Recall) 评估模型在多轮优化中的性能。

4. 实验结果与分析

文章通过多组实验，深入评估了所提出方法的有效性，并与现有方法进行了对比，主要结果如下：

结果 1: 排序损失微调优于回归微调

在 Table 1 中，作者对比了排序损失微调和回归微调 (MSE) 的性能。结果显示：

• 对于家族式模型 POET，排序损失微调在所有数据集大小和景观类型下均表现最佳。回归微调也是一个强基线，性能略优于最佳 ProteinNPT 配置。
• 对于掩码模型 ESM-1v 和 ESM-2，排序损失微调在绝大多数情况下显著优于回归微调，除了在 n=512 的多突变景观设置下。作者认为在 n=512 的多突变设置下，线性回归头可能更好地捕捉了上位性效应 (epistasis effects)。
• 随着训练数据量 n 的增加，两种微调方法之间的性能差距缩小。

结果 2: 排序损失微调优于基于冻结嵌入的模型

作者将最佳的排序损失微调方案与基于冻结嵌入的 ProteinNPT 基线模型进行比较。结果表明：

• POET 的排序损失微调在所有设置下都优于 ProteinNPT，尤其是在低数据量情况下，表明直接微调 POET 模型的似然函数对于最大限度地利用有限数据至关重要。
• 虽然 ESM-1v 和 ESM-2 通过 wt-marginal 排序策略微调后，在单突变景观上优于 ProteinNPT，但在多突变景观上表现稍逊，这可能是由于 wt-marginal 评分规则的表达能力有限。

结果 3: 排序损失微调能更好地泛化到未见位点

Table 2 评估了微调方法在测试集中未见突变位点的泛化能力。结果表明：

• 虽然在未见位点上的性能有所下降，但所有模型的似然函数直接进行排序损失微调，都展现出最佳的分布外 (out-of-distribution) 泛化能力。

结果 4: 掩码评分策略捕捉上位性效应

Table 3 比较了不同的掩码 PLM 评分策略在多突变 ProteinGym 景观上的性能。结果表明：

• "modulo" 掩码策略 (Johnson et al., 2024) 在 n=128 或 512 时优于其他策略，但需要更多的计算量。
• 高效的 wt-marginal 策略在计算效率和性能之间取得了良好的平衡。

结果 5: 集成方法提升性能

Table 4 和 Figure 1 展示了集成 PLM 模型在适应性预测和序列设计任务上的性能。结果表明：

• POET MSA 集成模型和掩码集成 ESM-2 模型都优于其对应的单一模型，证明了集成策略的有效性。
• POET 排序损失 MSA 集成模型在所有方法中表现最佳，尤其是在序列设计任务中，AUC 值显著高于其他方法。
• 在序列设计任务中，排序损失微调优于回归微调，集成策略进一步提升了整体性能。

5. 总结与评价

本文深入研究了基于似然函数微调蛋白质语言模型 (PLMs) 用于少样本适应性预测和序列设计的方法。通过扩展基于排序的损失函数，作者成功地将排序损失微调策略应用于掩码 PLM (ESM-1v, ESM-2) 和家族式自回归 PLM (POET) 模型，并取得了显著的性能提升。

文章的主要贡献和亮点包括:

• 系统性地比较了不同微调策略和不同类型 PLM 模型 在少样本蛋白质适应性预测任务上的性能，为领域研究人员提供了有价值的经验证据。
• 证明了排序损失微调方案的有效性，尤其是在低数据量情况下，相较于回归微调和基于冻结嵌入的方法，能够更好地利用 PLMs 的预训练知识。
• 提出了与排序损失微调兼容的集成策略，进一步提升了模型性能，并成功应用于多轮贝叶斯优化序列设计任务。
• 深入分析了模型在分布外泛化和捕捉上位性效应方面的能力，为理解 PLMs 的局限性和改进方向提供了 insights。

文章的局限性或未来方向:

• 虽然 "modulo" 掩码策略在多突变景观上表现出潜力，但其计算成本较高，未来可以探索更高效的上位性效应建模方法。
• 文章主要关注低数据量场景，未来可以研究如何将所提出的方法扩展到更大规模的数据集和更复杂的蛋白质工程任务。
• 可以进一步探索其他类型的集成策略，例如模型平均 (model averaging) 或模型蒸馏 (model distillation)，以提升集成模型的效率和泛化能力。

总而言之，这篇文章为蛋白质语言模型在少样本蛋白质工程领域的应用提供了重要的理论和实践指导，证明了基于似然函数和排序损失的微调方法，以及集成策略，是提升 PLMs 性能的有效途径，尤其是在数据受限的实际场景中。

7. Reference

Likelihood-based Finetuning of Protein Language Models for Few-shot Fitness Prediction and Design

代码链接: 文中没有明确给出代码链接，但提到了 ProteinNPT 的代码是公开可用的，可以参考 Notin et al. (2023b) 的论文或代码仓库。