5003.大模型-架构-AbAffinity-0.框架架构与使用

AbAffinity（harshitsinghsnu/AbAffinity）是面向抗体–抗原结合亲和力的仅序列（sequence-only）预测框架：以重链、轻链、抗原为三条独立但可互作的流，冻结 ESM-2（650M） 表征后，用**重链互补决定区（Complementarity-Determining Region，CDR）聚焦池化、轻重链自注意力、可学习融合门与门控交叉注意力（gated cross-attention）**估计 (\mathrm{p}K_{\mathrm{d}})。模型参数几乎只落在窄小的互作模块上，适合在缺少可靠复合物结构时做库规模筛选、零样本打分与少样本亲和力成熟微调。

段末注释：(\mathrm{p}K_{\mathrm{d}} = -\log_{10} K_{\mathrm{D}})，(K_{\mathrm{D}}) 越大则 (\mathrm{p}K_{\mathrm{d}}) 越小；CDR 为抗体可变区形成抗原接触的超变环；ESM-2 是 Meta 的蛋白质语言模型，输出每位氨基酸隐向量。

关联笔记：抗体亲和力预测工具调研 · 亲和力指标总览 · SPR/BLI 原理

官方资源：代码与权重 GitHub · MIT · Colab notebook 见仓库 README

---

1. 要解决什么问题

治疗性抗体发现、B 细胞受体（B-cell Receptor，BCR）组库分析与亲和力成熟都依赖对结合强弱的快速估计，然而：

1. 结构模型（对接、复合物精修）信息密度高，但依赖可靠三维复合物；库规模筛选时结构往往不可得或不可靠。
2. 序列模型吞吐高，却常把轻重链拼接成一条，或在早期就用简单融合把抗体与抗原混在一起，弱化了「哪条链怎么配对」「抗体与抗原是否兼容」的结构化建模。
3. 工业场景需要相对排序（谁更强）与部分定量 (\mathrm{p}K_{\mathrm{d}})，且顶点突变、单域纳米抗体（nanobody）与常规 IgG 配对抗体要同一套接口。

AbAffinity 的定位：不吃结构，只吃三条序列；把重链侧重到 CDR 环，用门控通路控制抗体–抗原信息交换，再用有界余弦头读出亲和力——ESM-2 冻结，只训互作头，训练成本低、迁移通道清晰。

---

2. 框架架构

官方终版模型类名为 MutualTriStreamStrong（AbAffinity/models/mutual_strong.py）。整体数据流如下图。

可按四段理解。

2.1 三流输入与冻结语言表征

输入为 light、heavy、antigen 三条氨基酸序列。上游用冻结 ESM-2（650M）（隐维 (d_{\mathrm{esm}}=1280)）为每条链生成残基级嵌入，再池化为向量：

流	池化方式	生物学意图
重链	All-CDR：在 CDR-H1 / H2 / H3 上聚焦池化	接触位点主要落在重链 CDR，尤其 CDR-H3
轻链	全链 mean pool	参与配对与部分副位（paratope），贡献更弥散
抗原	全链 mean pool	表位（epitope）位置先验弱于抗体 CDR，整链平均作为抗原指纹

编号优先 IMGT（经 ANARCI）；编号失败时改用基序（motif）回退。仓库自带的 data/*.pkl 缓存即在该配置下预计算，加载即可复现论文报告数字，无需再跑 ESM-2。

纳米抗体无重链：无有效轻链时，把重链嵌入复用到 light 槽位，同一架构覆盖常规配对抗体与单域抗体。

2.2 三路投影与轻重融合

三条池化向量先做独立 LayerNorm，再经各自的线性层投影到统一维度 (d=256)：

[
\mathbf{h} = f_H(\mathbf{e}_H),;
\mathbf{l} = f_L(\mathbf{e}_L),;
\mathbf{a} = f_A(\mathbf{e}_A)
]

随后只在抗体内部做一次双向交流：把 ([\mathbf{h},\mathbf{l}]) 堆成长度 2 的序列，过 多头自注意力（Multi-Head Self-Attention），得到更新后的 (\mathbf{h}’,\mathbf{l}’)。可学习 fusion_gate 以 (\mathrm{softmax}) 权重混合两条链：

[
\mathbf{w} = \mathrm{softmax}(W[\mathbf{h}’;\mathbf{l}’]),\quad
\mathbf{ab} = w_H\mathbf{h}’ + w_L\mathbf{l}’
]

消融实验表明：CDR 聚焦池化、轻重分解、真实抗原输入与门控互作各自贡献性能——早融合「一条拼接序列过全连接」的基线明显更弱。

2.3 门控双向交叉注意力

互作核心是可堆叠（默认 n_layers=2）的 GatedCrossAttention：抗体与抗原并行双向更新。

标准交叉注意力里 query 来自一方、key/value 来自另一方；AbAffinity 额外引入由 query 自身预测的 sigmoid 门：

[
\begin{aligned}
\mathbf{q},\mathbf{k},\mathbf{v} &= W_q\mathbf{x}q,; W_k\mathbf{x}{kv},; W_v\mathbf{x}_{kv}, \
\mathbf{u} &= \tanh(\mathbf{q}\odot\mathbf{k})\odot\mathbf{v}, \
\mathbf{g} &= \sigma(W_g\mathbf{x}_q), \
\mathbf{y} &= \mathrm{LN}\bigl(\mathbf{x}_q + W_o(\mathbf{u}\odot\mathbf{g})\bigr)
\end{aligned}
]

每一层并行两条通路：

• Ab → Ag：以融合抗体 (\mathbf{ab}) 为 query，当前抗原上下文为 key/value；
• Ag → Ab：以抗原为 query，融合抗体为 key/value。

门控使信息交换可「关小」——当某侧对当前亲和力预测贡献不大时不强制全量混入，这是相对朴素双塔或无门交叉注意力的关键差异。

2.4 双头余弦读出与 (\mathrm{p}K_{\mathrm{d}}) 标定

更新后的抗体 / 抗原向量分别进入 antibody_head / antigen_head（Linear → LN → GELU → Dropout），再 L2 归一化，输出余弦相似度 (s\in[-1,1])：

[
s = \cos(\hat{\mathbf{ab}},; \hat{\mathbf{ag}})
]

训练时把训练集 (\mathrm{p}K_{\mathrm{d}}) 线性缩放到 ([-1,1])：

[
t = 2\cdot\frac{y - y_{\min}}{y_{\max}-y_{\min}} - 1
]

以 (\mathrm{MSE}(s,t)) 优化；推理时用同一 (y_{\min},y_{\max}) 反变换得到预测 (\mathrm{p}K_{\mathrm{d}})。有界余弦头抑制表示坍缩（anti-collapse），使「相似度 → 亲和力」映射稳定、可解释为几何对齐强度。

段末注释：paratope / epitope 分别为抗体接触面与抗原被接触面；Integrated Gradients（IG） 用积分梯度归因，仓库 explain/ 用于把重要残基投回序列与三维结构。

---

3. 工作原理：端到端推演

以一对常规 VH–VL 抗体与抗原为例，推理路径可写成：

light / heavy / antigen 序列
        │ ESM-2（冻结）残基嵌入
        ▼
  轻/抗：mean pool　重：CDR-H1/H2/H3 pool
        │ 三路投影 → d=256
        ▼
  轻重自注意力 → fusion_gate → 抗体向量 ab
        │ 与抗原并行门控交叉注意力（可多层）
        ▼
  antibody_head / antigen_head → cos(ab, ag) ∈ [-1,1]
        │ 用训练 fold 的 (pKd_min, pKd_max) 反变换
        ▼
  预测 pKd（越大结合越强）

训练只更新投影、自注意力、融合门、门控交叉注意力与双头；ESM-2 不反传。因此：

• 预训练语义保留在三条流里；
• 「谁对谁重要、轻重占比、Ab–Ag 是否兼容」由窄小互作头学习；
• 零样本：权重固定，直接对 CSV 打分；
• 少样本：用少量带标本复合物微调互作头（验证集早停，独立测试集评估）。

可解释性侧，Integrated Gradients 对最终 (s)（或 (\mathrm{p}K_{\mathrm{d}})）做残基级归因， README 报告能高保真回收已知副位 / 表位残基；notebook 还可把高归因位点映射到三维结构（py3Dmol）上高亮。

---

4. 数据、评测与报告性能

4.1 输入格式

所有配对表与自定义 CSV 统一四列：

列	含义
light	轻链氨基酸序列（纳米抗体可空，运行时回退复用重链）
heavy	重链氨基酸序列
antigen	抗原氨基酸序列
Y	亲和力标签，单位 (\mathrm{p}K_{\mathrm{d}})

4.2 基准与划分

数据场景	用途	要点
SAaIntDB	主训 / 主报	10-fold CV；random（行级）与 cold（按 PDB_ID，禁同复合物泄漏）
SAbDab / AB-Bind / SKEMPI	MVSF-AB 外部基准	configs/benchmark/；可 SAbDab 全量训练、外推到 held-out 集合
预构建缓存	复现	data/allcdr_natural_650M.pkl、allcdr_mutation_650M.pkl、esm2_embeddings_saaintdb_650M.pkl 等

终版权重：model_weights/saaintdb_allcdr_random_bestfold.pt（SAaIntDB All-CDR，random 划分最佳 fold）。

4.3 SAaIntDB 上报告数字

十折交叉验证（仓库 README）：

指标	均值 ± 标准差
Pearson (r)	0.858 ± 0.006
Spearman (\rho)	0.844 ± 0.009
RMSE	0.694 ± 0.014

消融表明 CDR 池化、轻重分解、真实抗原与门控通路均不可省。外部变异图谱与基准上迁移有效，支撑结构不全时的零 / 少样本亲和力成熟筛查。绝对 (\mathrm{p}K_{\mathrm{d}}) 仍受同源复合物泄漏与批次标签噪声影响，cold 划分结果更接近配对面泛化能力。

---

5. 应用场景

场景	怎么用 AbAffinity	注意
序列库规模初筛	无结构、仅序列，批量零样本 (\mathrm{p}K_{\mathrm{d}}) / 排序	优先看相对排序，绝对值需 SPR/BLI 校准
亲和力成熟突变扫描	固定抗原，突变库走 few-shot 微调后打分	标签尽量同平台（都是 (\mathrm{p}K_{\mathrm{d}})），勿混 ELISA (EC_{50})
纳米抗体工程	无 light 时自动复用 heavy 槽	与常规 VH–VL 同一模型
副位 / 表位线索	IG 归因 + 可选结构映射	归因 ≠ 实验表位，需实验验证
BCR / 组库分析	对配对序列批量估结合倾向	缺乏真实抗原序列时不可硬套

在工具谱系中，AbAffinity 属于序列 + PLM + 结构感知互作头一类：不做 ΔΔG 物理能量（FoldX）、也不依赖全原子复合物图网络；更接近「可解释的三塔互作回归器」。选型时可与 抗体工具综述 中的排序任务、结构 ML 方法对照。

---

6. 安装与使用

6.1 环境

git clone https://github.com/harshitsinghsnu/AbAffinity.git
cd AbAffinity
conda env create -f environment.yml   # 与报告数字一致的 pin
conda activate AbAffinity
# CUDA 12.1 示例：
# pip install torch==2.2.1 --index-url https://download.pytorch.org/whl/cu121
pip install -e .

核心依赖：Python 3.10、PyTorch 2.2、transformers（ESM-2）、captum（IG）、biopython、ANARCI（IMGT 编号）、py3Dmol（notebook 三维展示）。

6.2 最快路径：自定义数据 notebook

仓库推荐入口：notebooks/AgAbGated_Custom_ZeroShot_FewShot_IG.ipynb（README 中的 Open in Colab）。

流程一键覆盖：

1. Zero-shot：加载 saaintdb_allcdr_random_bestfold.pt，对 CSV 预测并报告 Pearson / Spearman / RMSE；
2. Few-shot：用部分标签微调互作头，验证集早停，留出测试集；
3. IG 热图：残基级归因；
4. 结构映射：高归因残基叠到三维结构内联显示。

内置 60 个 SAaIntDB 样本可直接跑；换 DATA_CSV 或 Colab 上传即可换自家序列。ESM-2 嵌入步建议 GPU，纯 CPU 也能完成。

6.3 命令行复现训练 / 基准

SAaIntDB（多 seed）：

python -m AbAffinity.training.run_saaintdb_multiseed \
  --config configs/saaintdb/sa_ours_allcdr_random.yaml

python -m AbAffinity.training.run_saaintdb_multiseed \
  --config configs/saaintdb/sa_ours_allcdr_cold.yaml

默认 seed 42 114 144，结果写到 results/results_saaintdb/<name>/aggregated_summary.csv。

外部 MVSF-AB 数据集：

# SAbDab / AB-Bind / SKEMPI 十折 CV
python -m AbAffinity.training.run_multiseed \
  --config configs/benchmark/exp02_ours_allcdr_cv.yaml

# SAbDab 全量训练 → held-out 外推
python -m AbAffinity.training.run_multiseed \
  --config configs/benchmark/exp04_ours_allcdr_benchmark.yaml

# 跑全部 benchmark 配置
python -m AbAffinity.training.run_multiseed --all

仅当改动原始序列、需重建缓存时：

python -m AbAffinity.data_prep.precompute_embeddings_saaintdb
python -m AbAffinity.data_prep.precompute_allcdr_cache

官方默认可直接用 data/ 下已发布的 pickle，无需重算。

6.4 仓库布局（inker 入口）

AbAffinity/
├── AbAffinity/
│   ├── models/          # MutualTriStreamStrong 及变体
│   ├── utils/           # 数据、指标、训练循环
│   ├── training/        # CV / 多 seed / 消融 runner
│   ├── explain/         # IG 归因、门重要性
│   └── data_prep/       # ESM-2 / All-CDR 缓存预计算
├── configs/             # saaintdb / benchmark YAML
├── data/                # 表与预计算嵌入
├── model_weights/       # 终版权重
└── notebooks/           # 自定义 / 画图 / 可解释性

---

7. 局限与实践建议

1. 强依赖 (\mathrm{p}K_{\mathrm{d}}) 尺度一致性：与 ELISA (EC_{50})、混批次标签勿直接回归同一目标。
2. 随机划分可Optimistic：评估泛化请同时看 cold（按 PDB） 结果。
3. 抗原序列必须是真实抗原：消融表明真实抗原输入关键；用错误 IDs 当抗原会系统性偏置。
4. 默认立体化学之外的：没有三维几何——顺反立体、糖基化、多聚状态等不在模型内；结构侧信息应用 IG 做后验检查而非替代。
5. 长抗原 mean-pool 会稀释表位信号：超长抗原可考虑裁切已知结构域再输入。

实践上：零样本先排序 → 少量本靶点标点做 few-shot → IG 圈副位突变亮点 → SPR/BLI 验证 是一条可持续的闭环。

---

8. 小结

维度	要点
输入	light + heavy + antigen 序列；标签 (\mathrm{p}K_{\mathrm{d}})
表征	冻结 ESM-2（650M）；重链 All-CDR，轻 / 抗 mean pool
核心	三流投影 + 轻重自注意力与融合门 + 门控双向 Ab↔Ag 交叉注意力
读出	双头 L2 归一向量之间的余弦相似度，映射到 (\mathrm{p}K_{\mathrm{d}})
优势	无结构、可解释（IG）、纳秒抗体兼容、零 / 少样本迁移
适用	库筛、亲和力成熟初排、结构缺失时的序列友好流程

AbAffinity 把抗体亲和力问题收成「三流 PLM 指纹 + 窄小门控互作模块」：语义由 ESM-2 承担，配对与兼容性由可学习的两跳互作（轻重融合 → Ab–Ag 门控）承担。需要结构级物理直觉时仍应回 互作原理 与湿实验平台；需要算库规模序列分数时，本文给出的架构与命令即可作为落地起点。

---

参考文献与链接

1. AbAffinity 代码与文档：https://github.com/harshitsinghsnu/AbAffinity
2. ESM-2：Lin et al., Science (2023) — Language models of protein sequences at the scale of evolution
3. Captum（Integrated Gradients）：https://captum.ai/
4. ANARCI / IMGT 编号：Dunbar & Deane, Bioinformatics (2016)