ProteinMPNN：基于深度学习的蛋白质序列设计模型

ProteinMPNN：基于深度学习的蛋白质序列设计模型

一.ProteinMPNN 介绍 ProteinMPNN（Protein Message Passing Neural Network）是一种基于深度学习的蛋白质序列设计模型，核心目标是解决“逆向折叠问题”（inverse folding problem），即根据给定的蛋白质三维结构，设计出能够折叠成该结构的氨基酸序列。ProteinMPNN 在计算和实验测试中都有出色的性能表现，不同位置的氨基酸序列可以在单链或多链之间偶联，从而广泛的应用于当前蛋白质设计上。ProteinMPNN 不仅在天然蛋白质序列恢复率上面性能要高于传统的 Rosetta 方法，并且可以恢复先前设计失败的蛋白质。通过前沿 AI 技术突破科学研究的效率瓶颈，对于蛋白质工程、药物设计、酶设计等领域有极其重要的意义。

二.整体架构

图 1: ProteinMPNN 的基本架构

1.编码器

1.1 图构建将蛋白质结构表示为图结构，图的节点代表氨基酸残基，边代表残基之间的空间或序列关系（如距离、接触、氢键等）。

1.2 节点嵌入使用可学习的嵌入层将每个残基的类型（氨基酸种类）、位置信息编码为初始节点特征。

是节点 v 在第 l 层的嵌入向量。

是节点 v 的邻居节点集合，

是从节点 u 到节点 v 的消息。

1.3 边嵌入对边进行编码，边的初始特征包括残基之间对距离、角度的几何变换，或残基间的相互作用特征。

是边（u, v）在第 l 层的嵌入向量，

分别是节点 u 和节点 v 的嵌入向量。

2.消息传递网络

模型通过多层堆叠的消息传递和节点更新操作迭代更新节点状态

2.1 多层图卷积在每一层图卷积中，节点 v 接收邻居节点 u 的消息 Muv,由边特征 euv 和节点特征 hu 计算得到：

其中

是第 l 层的可学习函数（如 MLP），

包含距离、角度等几何信息。

2.2 SE3-等变注意力机制

2.2.1 注意力分数计算计算节点 v 对邻居 u 的注意力分数

，结合系欸但特征和边特征：

其中 Wa,Wb,Wc 是可学习权重矩阵，

是层间参数，softmax 确保归一化。

2.2.2 消息聚合（SE3-等变）消息传递需满足 SE3 群的对称性，通常通过张量积实现：

其中

，

是可学习的变换矩阵，k 是 SE3 的表示维度。

3.解码器

3.1 结构预测头用于预测三位坐标，二面角或距离矩阵。

是激活函数，

是坐标回归权重。

3.2 序列设计头

3.2.1 自回归生成（交叉熵损失）逐个生成氨基酸序列，条件于当前节点特征和已生成序列：

是时间 t 的氨基酸类型，

为权重矩阵。

3.2.2 损失函数（序列恢复）交叉熵损失对比预测序列与天然序列：

4.多任务学习

同时优化结构预测（如坐标误差）和序列设计（如恢复天然序列）任务，通过联合损失函数平衡不同目标。

三.实验

设备：

Atlas 800T A2

组件版本：

hdk：24.1.RC3

cann：8.0.RC3

python：3.10

torch：2.1.0

torch-npu：2.1.0.post8

1. 环境准备

官方文档对依赖包没有版本要求，仅需保证 python 的版本高于等于 3.0 即可。

1.1 新建 conda 环境 conda create --name mlfold python=3.10

1.2 安装 torch_npu 和一些必要的依赖包下载 torch_npuwget https://gitee.com/ascend/pytorch/releases/download/v6.0.rc3-pytorch2.1.0/torch_npu-2.1.0.post8-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl --no-check-certificate

安装 pip3 install torch_npu-2.1.0.post8-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl

注：该环境已经提前安装过了，所以不会提示安装成功

2.数据集下载

2.1 下载原始数据集 wget https://files.ipd.uw.edu/pub/training_sets/pdb_2021aug02.tar.gz

注：原始 pdb_2021aug02 数据集压缩包大小为 17GB，解压之后为 92GB，须预留足够的空间。下载之后使用 tar -xzf pdb_2021aug02.tar.gz 进行解压。

2.2 下载用于测试的数据集样本如果没有足够的磁盘空间，可下载该数据集，压缩包为 48MB，解压之后为 255MB。wget https://files.ipd.uw.edu/pub/training_sets/pdb_2021aug02_sample.tar.gz

3.推理

ProteinMPNN 在 example 目录下面有很多样例，均可用来验证推理。

在 training 目录下也有 test_inference.sh 用于推理验证。

3.1 激活 conda 环境 conda activate mlfold

3.2 修改推理脚本 Slurm 参数、用于输入的.PDB 文件、权重路径、权重名称

3.3 修改 protein_mpnn_run.py

3.4 执行推理

3.5 查看推理结果

4.训练

4.1 修改 training.py 文件

注：training.py 脚本内，使用的 cuda api 不需要手动修改。torch_npu 会自动将 cuda api 更改为 torch_npu 的 api。

例：

4.2 修改 sh 脚本

4.3 source cann

4.4 验证 npu 环境是否可用

4.5 开始训练

在 training 目录下面执行 bash submit_exp_020.sh

查看 npu 状态，显示 0 卡正在被进程占用

注：可以通过环境变量 export ASCEND_RT_VISIBLE_DEVICES 来指定用哪张卡训练。

4.6 训练完成

训练完成之后，训练日志和权重会存放在 exp_020 目录下。