ProteinMPNN：概述与训练框架

发表于 | 更新于 | 分类于 ， | 阅读 次
ProteinMPNN 架构与训练全景：骨架图输入、顺序无关自回归目标、Encoder-Decoder MPNN（EncLayer / DecLayer 伪代码）与约 1.68M 参数组成；与 RFdiffusion pipeline 及 MPNN 抽象对照。

ProteinMPNN（Dauparas et al., Science 2022）是 蛋白质逆折叠（inverse folding） 的代表模型：给定骨架三维结构（backbone coordinates），预测氨基酸序列，使序列在该构象下稳定折叠。与 AlphaFold2「序列 → 结构」相反，ProteinMPNN 服务 de novo design pipeline 中「结构已定 → 填序列」环节——典型组合为 RFdiffusion / Chroma 生成骨架 → ProteinMPNN 设计序列 → AlphaFold2 自洽筛选（见 Diffusion-5）。

段末注释：逆折叠 指由目标构象反推氨基酸序列；de novo design 指从头设计而非天然序列微调。

前置阅读：MPNN-0 消息传递神经网络、范式综述 §GNN
对照：Diffusion-Example-1 序列离散扩散（纯序列生成范式）

插图：5003.大模型-架构-ProteinMPNN-0.概述/proteinmpnn-fig*.png
官方代码：github.com/dauparas/ProteinMPNN

1. 任务定义与在 design pipeline 中的位置

步骤 模型/工具 输入 → 输出
1. 骨架生成 RFdiffusion、Chroma 等 条件 → 3D 骨架
2. 序列设计 ProteinMPNN 骨架 → 氨基酸序列
3. 验证 AlphaFold2、ESMFold 序列 → 预测结构，算 scRMSD / pLDDT

ProteinMPNN 不预测结构，只建模 (p(\mathrm{seq}\mid\mathrm{backbone}))；结构质量由上游扩散/优化与下游折叠器负责。

2. 结构框架总览

ProteinMPNN = 稀疏 kNN 图 + Encoder MPNN（仅看骨架）+ Decoder MPNN（骨架 + 部分序列上下文）+ 顺序无关自回归（order-agnostic autoregressive） 解码。

图 1　ProteinMPNN 整体架构：kNN 图 → Encoder → 部分序列 + 图特征 → Decoder → 氨基酸 logits

2.1 图构建

  • 节点：每个残基一个节点（以 Cα 为中心建 k 近邻图）。
  • 边：每个节点连 (k) 个空间最近邻（官方默认 (k=48)，按 Cα–Cα 距离）。
  • 多链：全复合物所有链的残基置于同一图；边特征含同链 / 异链标识。

2.2 Encoder（骨架编码器）

  • 3 层消息传递（堆叠 3 个 EncLayer，不是 MLP 内部的 3 个线性层），hidden_dim = 128，含 边更新（edge update）（非仅节点更新）。
  • 输入：不依赖氨基酸类型（或仅用 mask 占位），主要编码几何：
    • N、Cα、C、O 与虚拟 Cβ 之间的距离（RBF 展开）；
    • 相对 Cα–Cα–Cα 局部坐标系方向/旋转；
    • 骨架二面角等（早期 baseline 以二面角为主，最终版以原子间距离为主，恢复率更高）。
  • 输出：每个残基节点特征 (\mathbf{h}i)、边特征 (\mathbf{e}{ij})（供 Decoder 复用）。
  • 实现要点：节点 (\mathbf{h}_i) 从零向量起步，几何全部走边；层内无因果掩码（全图双向）；先更新节点，再用新节点更新边。对应官方 EncLayer。

一层 Encoder（论文 SI + protein_mpnn_utils.py 的 EncLayer）：

1
2
3
4
5
6
7
8
9
10
11
12
# encoder_layer: 同时更新节点 V 与边 E（官方 EncLayer；全图并行，不是按 i 串行）
# 输入: V[L,128] 节点；E[L,k,128] 边；N(i)=i 的 kNN
# 输出: 更新后的 V, E
# 处理: 全体边算消息 → 聚合更新全体 V → 再用新 V 写回全体边
def encoder_layer(V, E):
    M_ij = MLP3(concat(V_i, V_j, E_ij))            # 对所有边并行；代码拼接为 [V_i, E_ij, V_j]
    dV_i = sum_j(M_ij) / 30                        # 论文 SI 写 Sum；代码 scale=30
    V    = LayerNorm(V + Dropout(dV))
    V    = LayerNorm(V + Dropout(FFN(V)))          # 128→512→128，GELU
    dE_ij = MLP3(concat(V_i, V_j, E_ij))           # 边更新用的是新 V；Decoder 没有这一步
    E    = LayerNorm(E + Dropout(dE))
    return V, E

堆叠 3 个 EncLayer（超参 num_encoder_layers=3）：V ← 0，E ← W_e(几何边特征)，再 for _ in 3: V, E = encoder_layer(V, E)。每个 EncLayer 内部另有自己的 MLP3 / FFN，那是层内算子，不是「Encoder 三层」。

段末注释：MLP3 不是官方类名，只表示「三层线性」的消息多层感知机（Multi-Layer Perceptron，MLP）：(\mathrm{MLP3}(x)=W_3,\mathrm{GELU}(W_2,\mathrm{GELU}(W_1 x)))，对应代码 W3(GELU(W2(GELU(W1(h_EV)))))。Encoder 的 (W_1) 是 (384\to 128)（拼接 ([\mathbf{v}i,\mathbf{e}{ij},\mathbf{v}_j])），其后两层 (128\to 128)。FFN（Position-wise FeedForward）是另一条支路：(128\to 512\to 128)，GELU，只作用在节点上。LayerNorm 为层归一化。

2.3 Decoder（序列解码器）

  • 3 层 MPNN（堆叠 3 个 DecLayer，num_decoder_layers=3），hidden 128。
  • 输入：Encoder 输出的 (\mathbf{h}i,\mathbf{e}{ij}) + 已部分解码的氨基酸序列（one-hot 或 embedding）。
  • 输出：每个位置 20 种标准氨基酸（+ 掩码 token）的 logits。
  • 顺序无关自回归：训练时解码顺序 (\pi) 从全排列中随机采样；推理时可任意顺序，先解码固定 motif、再设计侧翼（Fig. 1B）。
  • 实现要点：不更新边；把邻居氨基酸 (\mathbf{s}j) 拼进边，用因果掩码 (\mathrm{mask}{ij}) 挡住「尚未解码」的序列；未解码邻居只看 Encoder 几何。对应官方 DecLayer。

一层 Decoder（论文 SI + DecLayer）：

1
2
3
4
5
6
7
8
9
10
11
12
# decoder_layer: 只更新节点；边拼接邻居氨基酸（官方 DecLayer；全图并行）
# 输入: V；E（Encoder 边，本层不改）；S 部分序列嵌入；mask_ij=1 当 j 已在 i 之前解码
# 输出: 更新后的 V
# 处理: 未来位置的 S_j 置零，消息仍含 V_j 与几何边
def decoder_layer(V, E, S, mask):
    S_ctx    = where(mask_ij, S_j, 0)              # 因果：未来位置无序列上下文
    E_ij_dec = concat(E_ij, S_ctx)
    M_ij     = MLP3(concat(V_i, V_j, E_ij_dec))
    dV_i     = sum_j(M_ij) / 30
    V        = LayerNorm(V + Dropout(dV))
    V        = LayerNorm(V + Dropout(FFN(V)))
    return V

训练一次前向（teacher forcing，全位置并行；官方 ProteinMPNN.forward）：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
# proteinmpnn_train_forward
# 输入: 骨架坐标 X；标签序列 S；随机解码顺序 π
# 输出: 每残基 21 类 logits
# 处理: Encoder 无掩码编码几何 → 按 π 造因果 mask → Decoder 并行读部分序列
def proteinmpnn_train_forward(X, S, pi):
    E, E_idx = ProteinFeatures(X)                 # Cα-kNN + 25 组原子对 RBF + 相对位置(±32, 同/异链)
    V = zeros(L, 128)                             # 几何全在边上
    E = W_e(E)
    for _ in range(3):
        V, E = encoder_layer(V, E)                # 无掩码

    S_emb  = Embed(S)                             # 21 → 128；标签序列（teacher forcing）
    mask_ij = 1 if pi_inv[j] < pi_inv[i] else 0   # j 在 π 中排在 i 之前
    # 官方写法：已解码邻居用 [E_ij, S_j, V_j]，未解码用 [E_ij, 0, V_j]
    for _ in range(3):
        V = decoder_layer(V, E, S_emb, mask)
    return Linear(V)                              # logits: 21 类

推理时 Encoder 仍算一次；Decoder 按 (\pi) 逐步写回已采样氨基酸（或 tied 组同时采样），不能再用整条标签 (S) 并行。固定 motif 时把这些位置排进 (\pi) 的前缀，使其作为后续位置的 (S_j)。

Encoder Decoder
看什么 仅骨架几何 几何 + 已解码氨基酸
节点初值 (\mathbf{0}) Encoder 的 (\mathbf{h})
边 更新（第 2 个 MLP3） 不更新；只拼 (S_j) 或 (\mathbf{0})
掩码 仅缺失残基 因果 / 顺序无关 AR
官方类 EncLayer(128, 256) DecLayer(128, 384)
一层输入拼接 ([\mathbf{v}i,\mathbf{e}{ij},\mathbf{v}_j]) ([\mathbf{v}i,\mathbf{e}{ij},\mathbf{s}_j\text{ or }0,\mathbf{v}_j])

段末注释：teacher forcing 指训练时用真实标签前缀而非模型自己的采样；官方用排列 (\pi) 构造下三角掩码，从而一次前向并行算完所有位置的条件分布。

2.4 对称 / 多链 tying

  • Tied positions：指定多个位置（可跨链）共享同一 logits（对数概率平均后采样），用于同源寡聚体、重复蛋白、环对称 design（Fig. 1C）。

3. 训练输入信息（详细）

图 2　ProteinMPNN 训练输入与监督目标

3.1 数据来源与过滤

项目 设置（论文 / 官方 training.py）
数据库 PDB（训练集快照 2021-08-02）
方法 X-ray 或 cryoEM
分辨率 ≤ 3.5 Å（rescut=3.5）
组装 生物组装（biological assembly），多链复合物
划分 序列聚类 30% identity；CATH 80/10/10（单体开发实验）；valid_clusters / test_clusters

官方 .pt 每条链字段（training README）：

字段 shape / 类型 含义
seq string 氨基酸序列（监督标签）
xyz [L, 14, 3] 原子坐标（骨架 + 侧链槽位）
mask [L, 14] bool 原子是否存在于 PDB
bfac / occ [L, 14] 温度因子、占有率

3.2 几何特征（模型实际「看到」的输入）

特征类 内容 作用
距离 RBF N、Cα、C、O、虚拟 Cβ 间距离（残基对 / 局部） 主归纳偏置；替代纯二面角更优（论文 Table 1 Exp.1）
局部取向 Cα 邻域刚体帧相对旋转 编码侧链朝向、氢键几何
序列位置 残基索引差 截断 ±32；同链 / 异链 bit 多链与界面设计
kNN 拓扑 Cα 距离 Top-(k) 邻接 稀疏 MPNN 计算域

3.3 训练时动态输入

机制 说明
随机解码顺序 (\pi) 每步从排列中采样；已解码位置提供序列上下文，未解码为 mask
骨架噪声 对坐标加高斯扰动（官方默认 backbone_noise=0.2 Å 量级；发布权重有 0.02 / 0.10 / 0.20 Å 三档）
Dropout 默认 0.1
Batch 按 token 数计，默认 10,000 tokens/batch（多蛋白拼 batch）

4. 训练目标

4.1 顺序无关自回归交叉熵

对解码顺序 (\pi=(\pi_1,\ldots,\pi_L))，在第 (s) 步预测位置 (\pi_s) 的氨基酸 (a_{\pi_s})：

[
\mathcal{L} = -\frac{1}{N_{\mathrm{norm}}}\sum_{s=1}^{L} \log p_\theta\big(a_{\pi_s} ,\big|, \mathrm{backbone},, a_{\pi_1},\ldots,a_{\pi_{s-1}}\big)
]

  • (p_\theta)：Decoder 输出的 21 类 softmax（20 AA + mask/unknown）。
  • (N_{\mathrm{norm}})：实现中按 batch 总 token 归一（论文报告 perplexity = (\exp()平均 NLL())）。
  • 无结构损失：不对坐标回归；结构信息只经 Encoder 进入条件。

段末注释：Perplexity 为 (\exp()交叉熵())，越低表示序列预测越 confident 且与标签一致。

4.2 评估指标（非训练 loss）

指标 含义
Sequence recovery 预测序列与天然序列相同的残基比例（%）
Perplexity held-out 集上的 (\exp(\mathrm{NLL}))
AlphaFold consistency 设计序列单序列折叠与输入骨架的 TM-score / scRMSD

论文最终 ProteinMPNN（Exp.3+4，随机顺序）：PDB 测试集 recovery ≈ 50.8%（无噪训练）/ 47.9%（训练时 0.02 Å 噪）；相对 baseline 41.2% 显著提升。

4.3 优化与 schedule

超参 典型值
优化器 Adam
Epoch 150–200（发布 v_48_* 为 150 epoch）
混合精度 默认开启
梯度裁剪 可选（gradient_norm）

5. 参数组成与模型变体

图 3　ProteinMPNN 参数模块组成（约 1.68M）

5.1 默认架构超参（与 training.py 一致）

超参数 默认值 说明
hidden_dim 128 节点/边隐层宽度
num_encoder_layers 3 骨架 Encoder MPNN 层数
num_decoder_layers 3 序列 Decoder MPNN 层数
num_neighbors 48 kNN 邻居数 (k)
dropout 0.1
参数量 ≈ 1.68 M 论文 Table 1 Exp.3/4（含边更新）

5.2 参数模块分解（数量级）

模块 角色 参数量级
距离 RBF + 边/节点嵌入 MLP 几何标量 → 向量 ~10⁴–10⁵
Encoder（3× MPNN + edge update） 骨架 → (\mathbf{h},\mathbf{e}) ~45%
Decoder（3× MPNN） 部分序列 + 图 → logits ~45%
氨基酸嵌入 + 输出线性层 21 类分类头 ~10⁴

增大 hidden_dim 主要降低训练 loss，验证 recovery 增益边际（论文：>128 收益有限）。

5.3 官方发布权重

权重文件 训练骨架噪声 (\sigma) 适用场景
v_48_002.pt 0.02 Å 近天然、高精度骨架
v_48_010.pt 0.10 Å 中等
v_48_020.pt 0.20 Å hallucinated / 低精度 骨架（RFdiffusion 常用）

推理时可调 temperature 控制序列多样性；固定位置设 fixed_positions 做 motif scaffolding。

6. 推理与训练的关键差异

训练 推理
序列上下文 随机顺序 (\pi) 下已知前缀 可指定顺序（先固定、后设计）
骨架 真实 PDB + 可选噪声 设计模型输出的 idealized 骨架
输出 CE loss 自回归采样 + optional tying
多链 生物组装图 界面 / binder design 同一套 API

7. 与 MPNN 抽象及 LLM 的对照

MPNN-0 抽象 ProteinMPNN
图 ((V,E)) kNN on Cα，(k=48)
消息 (\mathrm{MSG}(h_u,h_v,e_{uv})) 几何边特征 + 边更新
读出 图级 pool 逐残基 氨基酸分类
序列性 无 inherent 顺序 顺序无关 AR

与 Transformer LLM 不同：ProteinMPNN 不做自回归语言建模；结构先验在稀疏 3D 图上，序列仅 ~1.7M 参数即可高 recovery——归纳偏置来自 几何距离 + 局部 MPNN 而非 billion-scale 预训练。

8. 小结

维度 ProteinMPNN 要点
输入 骨架 N/Cα/C/O（+ 虚拟 Cβ）坐标 → kNN 图 + RBF 距离/取向；多链 biounit
目标 顺序无关自回归 交叉熵，预测天然序列
结构 3+3 层 MPNN Encoder–Decoder，hidden=128，~1.68M 参数
工程 骨架噪声增强、tied positions、固定残基 inpainting
生态 RFdiffusion 之后的事实标准序列设计模块

本篇索引

主题 章节 延伸阅读
MPNN 消息传递 §2 MPNN-0
Encoder / Decoder 层实现 §2.2–2.3 官方 EncLayer / DecLayer（protein_mpnn_utils.py）
训练数据格式 §3.1 官方 training README
损失与 recovery §4 论文 Table 1
design pipeline §1 Diffusion-5
序列扩散替代 §7 Diffusion-Example-1

参考与延伸阅读

  • Dauparas et al., Robust deep learning–based protein sequence design using ProteinMPNN, Science 378, 49–56 (2022). doi:10.1126/science.add2187
  • Jumper et al., AlphaFold2（结构预测对照）.
  • Watson et al., RFdiffusion（上游骨架生成）.
  • 官方实现：dauparas/ProteinMPNN.
-------------本文结束感谢您的阅读-------------