HOOOS

如何用 ESM-2 进行抗体-抗原结合亲和力预测？从零样本表征到微调实操

0 5 孤波 ESM-2抗体药物研发亲和力预测
Apple

在 AI 辅助抗体药物研发（AIDD）中，评估抗体与抗原之间的结合亲和力（Affinity）是核心环节。Meta 团队开源的 ESM-2 作为目前最强大的蛋白质语言模型之一，凭借其在海量无标注蛋白质序列上学习到的进化和物理化学规律，成为了这一任务的首选基座。

然而，ESM-2 本质上是一个单序列遮盖语言模型（Masked Language Model, MLM）。如何将一个预测“填空题”的模型，转化为能够精确预测抗体-抗原相互作用（CPI/API）的工具？

零样本（Zero-shot）预测的物理逻辑与实现

所谓的“零样本预测”，是指不使用任何亲和力定量数据（如 $K_d, K_i, \text{IC}_{50}$）进行微调，直接利用 ESM-2 的输出指标来评估亲和力的相对改变（$\Delta \Delta G$ 或 $\Delta pK_d$）。

这种方法主要用于**抗体亲和力成熟（Affinity Maturation）**阶段的突变筛选。

1. 核心原理：伪对数似然差值（Pseudo-Log-Likelihood Difference）

语言模型对一个序列输出的概率，反映了该序列在进化上的“合理性”。如果我们对抗体（特别是 CDR 区）进行单点或多点突变，可以用 ESM-2 计算突变体相比于野毒株/野生型的概率变化。

对于抗体-抗原复合物，最直接的零样本 Trick 是拼接序列法（Concatenation Linker）。将抗体可变区序列（如 scFv 或 Heavy/Light chain）与抗原靶标序列用一段特殊的 Linker（例如数十个甘氨酸 G，或者直接用 <span> 占位符）拼接成一条长单链送入 ESM-2。

定义突变前后的伪对数似然变化（$\Delta \text{PLL}$）：

$$\Delta \text{PLL} = \sum_{i \in \text{mut}} \log P(x_i = \text{mut}i \mid \mathbf{X}{\backslash i}) - \sum_{i \in \text{mut}} \log P(x_i = \text{wt}i \mid \mathbf{X}{\backslash i})$$

其中，$\mathbf{X}$ 是抗体-抗原拼接后的整链。$\Delta \text{PLL}$ 越高，代表突变后系统整体的能量基态越稳定，在统计学上与亲和力提高呈正相关。

2. 零样本预测的 Python 代码实现

以下是使用 Hugging Face transformers 库计算抗体突变后 Pseudo-Log-Likelihood 的核心逻辑：

import torch
from transformers import AutoTokenizer, EsmForMaskedLM

# 初始化 ESM-2 模型（推荐使用 650M 或 3B 参数版本以获得高精度）
model_name = "facebook/esm2_t33_650M_UR50D"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = EsmForMaskedLM.from_pretrained(model_name).eval().cuda()

def get_sequence_loss(sequence, mutate_indices=None):
    """
    计算特定位置处的 Pseudo-log-likelihood.
    如果 mutate_indices 为 None，则计算整条序列的平均 PLL。
    """
    inputs = tokenizer(sequence, return_tensors="pt")
    input_ids = inputs["input_ids"].cuda()
    attention_mask = inputs["attention_mask"].cuda()
    
    with torch.no_grad():
        outputs = model(input_ids, attention_mask=attention_mask)
        logits = outputs.logits # [1, seq_len, alphabet_size]
        
    log_probs = torch.log_softmax(logits, dim=-1)
    
    # 提取实际氨基酸的对数概率
    # input_ids: [1, seq_len]，除去 [CLS] 和 [SEP]
    target_log_probs = log_probs.gather(dim=-1, index=input_ids.unsqueeze(-1)).squeeze(-1)
    
    if mutate_indices is not None:
        # 仅累加突变位点的 log probability
        # 注意：此处索引需对应 tokenizer 编码后的 token 索引
        score = target_log_probs[0, mutate_indices].sum().item()
    else:
        score = target_log_probs[0, 1:-1].mean().item()
    return score

# 示例：抗体重链可变区与抗原 Epitope 拼接（用G接头）
wild_type = "EVQLVESGGGLV...GGGGSGGGGS...KVECRYV" # 抗体 + Linker + 抗原
mutant_type = "EVQLVESGGALV...GGGGSGGGGS...KVECRYV" # 第 10 位 V -> A

# 计算突变点在 tokenizer 后的真实索引 (假设为 10)
mut_idx = [10] 

score_wt = get_sequence_loss(wild_type, mut_idx)
score_mut = get_sequence_loss(mutant_type, mut_idx)

delta_pll = score_mut - score_wt
print(f"Relative affinity score (Delta PLL): {delta_pll:.4f}")

避坑指南： 零样本预测抗体亲和力高度依赖于抗原的存在。如果仅仅输入抗体单链，模型评估的其实是**抗体自身的折叠稳定性（Stability）**而非结合亲和力。因此，拼接抗原序列是必要的折中方案，尽管这超出了 ESM-2 预训练时单链无接头的分布。

适配抗体-抗原亲和力预测的微调（Fine-tuning）策略

当手里拥有几十到几千个具有确切亲和力常数（如 $K_d$ 值）的抗体-抗原湿实验数据时，微调（Fine-tuning）能够极大提升模型的绝对定量预测精度。

根据算力和数据规模，有以下三种主流微调架构：

策略一：双塔（Bi-Encoder）表征融合 + 下游回归机（轻量高效）

这是工业界最高效的基线（Baseline）方案。

[抗体序列] ---> ESM-2 (Frozen) ---> [抗体 Pooling 向量] \
                                                          (拼接/相减/Hadamard 积) ---> MLP / XGBoost ---> Log(Kd)
[抗原序列] ---> ESM-2 (Frozen) ---> [抗原 Pooling 向量] /
  • 操作方法：
    1. 冻结（Freeze）ESM-2 全部的 Transformer 层，不参与梯度更新。
    2. 将抗体序列和抗原序列分别送入 ESM-2，提取最后一层的表征。
    3. 抗体端通常只提取 CDR 区（CDR-H1/2/3, CDR-L1/2/3） 对应 Token 的 Mean-pooling 向量；抗原端提取 Epitope（表位） 区域的 Mean-pooling 向量。
    4. 将两个特征向量进行组合：$[v_{ab}, v_{ag}, |v_{ab} - v_{ag}|, v_{ab} \odot v_{ag}]$。
    5. 训练一个轻量级的 MLP（多层感知机）或 XGBoost 预测 $\log(K_d)$。
  • 适用场景：数据量极少（$< 500$ 条），或者算力极度受限（只有单张显卡）。此方案完全避免了过拟合。

策略二：参数高效微调（PEFT / LoRA）

当数据量达到几千条时，直接微调 650M 以上参数的 ESM-2 会面临显存崩溃和严重的灾难性遗忘。此时引入 LoRA (Low-Rank Adaptation) 是最佳策略。

  • 微调配置：
    • 将 LoRA 注入到 ESM-2 自注意力模块的 query 和 value 投影矩阵中。
    • 使用双塔架构或单塔拼接架构。对于亲和力预测，单塔（将抗体与抗原直接拼接后送入模型）+ 全注意力机制效果最好，因为自注意力机制能够直接学习抗体残基与抗原残基之间的两两相互作用。
  • 代码实操 (PEFT 整合)：
from peft import LoraConfig, get_peft_model

# 针对 ESM-2 的 LoRA 配置
config = LoraConfig(
    r=8, 
    lora_alpha=16, 
    target_modules=["query", "value"], # ESM-2 的 Attention 映射命名
    lora_dropout=0.1, 
    bias="none",
    modules_to_save=["classifier"] # 顶部的回归头
)

model = get_peft_model(raw_esm_model, config)
model.print_trainable_parameters()

策略三：基于对比学习（Contrastive Learning）的表征对齐

如果你的任务不是回归具体的 $K_d$ 值，而是进行高通量虚拟筛选（在百万级抗体库中找出能结合特定抗原的抗体），对比学习是上限最高的设计。

  • 架构设计：
    • 构建双塔：抗体编码器 $E_{ab}$ 和抗原编码器 $E_{ag}$（均由 ESM-2 初始化并应用 LoRA）。
    • 通过投影层将两者的表征投影到同一低维流形空间（例如 256 维）。
    • 使用 InfoNCE Loss，拉近有结合活性的抗体-抗原对（Positive pairs）在空间中的距离，推远不结合的配对（Negative pairs）。
  • 优势：在测试阶段，你可以预先计算抗原的 Embedding，并对整个抗体库进行一次性 Embedding 化，通过快速的**余弦相似度（Cosine Similarity）**计算检索出潜在抗体，检索速度比逐个进行序列拼接预测快 4 个数量级。

避坑指南：抗体亲和力预测的独特挑战

在处理抗体数据时，直接套用通用的蛋白质表征微调方法往往会遭遇失败，必须注意以下几点：

  1. 致命的数据泄漏（Data Leakage）：基于 CDR 聚类划分数据集
    不能简单地使用随机划分（Random Split）来划分训练集和测试集。抗体的骨架区（Framework）高度相似，如果随机划分，会导致测试集中大量抗体的 CDR 序列与训练集极度相似。必须采用 CDR-H3 序列同源性聚类划分（如使用 MMseqs2 保持测试集与训练集 CDR 序列恒等性 $< 80%$），否则评估出来的模型精度会虚高。
  2. 长度外推问题
    抗原序列通常很长。如果将抗原全长与抗体直接拼接，极易超出 ESM-2 的 1024 默认上下文窗口。建议：切片（Crop）抗原。 只保留已知的表位区域（Epitope）及周边 10-15 个氨基酸残基，抛弃无关的抗原结构域。
  3. 坐标缺失时的结构辅助
    纯序列微调的上限较低。如果在微调时能引入结构特征，例如用 ESM-Fold 快速预测抗体的 3D 结构，提取残基的物理距离矩阵（Distance Matrix），作为偏置（Bias）注入到 ESM-2 的 Attention Map 中，能有效弥补一维序列模型对空间构象不敏感的缺陷。

点评评价

captcha
健康