一 ProteinMPNN介绍
ProteinMPNN（Protein Message Passing Neural Network）是一种基于深度学习的蛋白质序列设计模型，核心目标是解决“逆向折叠问题”（inverse folding problem），即根据给定的蛋白质三维结构，设计出能够折叠成该结构的氨基酸序列。ProteinMPNN在计算和实验测试中都有出色的性能表现，不同位置的氨基酸序列可以在单链或多链之间偶联，从而广泛的应用于当前蛋白质设计上。ProteinMPNN不仅在天然蛋白质序列恢复率上面性能要高于传统的Rosetta方法，并且可以恢复先前设计失败的蛋白质。通过前沿AI技术突破科学研究的效率瓶颈，对于蛋白质工程、药物设计、酶设计等领域有极其重要的意义。
二 整体架构
图1: ProteinMPNN的基本架构
2.1 编码器
2.1.1 图构建
将蛋白质结构表示为图结构，图的节点代表氨基酸残基，边代表残基之间的空间或序列关系（如距离、接触、氢键等）。
2.1.2 节点嵌入
使用可学习的嵌入层将每个残基的类型（氨基酸种类）、位置信息编码为初始节点特征。
\( \mathcal{N}(v) \)是节点v在第l层的嵌入向量。
\( m_{u v}^{(l)} \) 是节点v的邻居节点集合，是从节点u到节点v的消息。
2.1.3 边嵌入
对边进行编码，边的初始特征包括残基之间对距离、角度的几何变换，或残基间的相互作用特征。
\( e_{u v}^{(l)} \)是边（u, v）在第l层的嵌入向量，
\( h_u^{(l)}, h_v^{(l)} \)分别是节点u和节点v的嵌入向量。
2.2 消息传递网络
模型通过多层堆叠的消息传递和节点更新操作迭代更新节点状态
2.2.1 多层图卷积
在每一层图卷积中，节点v接收邻居节点u的消息Muv,由边特征euv和节点特征hu计算得到：
其中，\( \phi^{(l)} \) 是第l层的可学习函数（如MLP），\( e_{u v} \)包含距离、角度等几何信息。
2.2.2 SE3-等变注意力机制
注意力分数计算
计算节点v对邻居u的注意力分数\( e_{v u} \)，结合系欸但特征和边特征：
其中Wa,Wb,Wc是可学习权重矩阵，\( \beta^{(l)} \) ，是层间参数，softmax确保归一化。
消息聚合（SE3-等变）
消息传递需满足SE3群的对称性，通常通过张量积实现：
其中，\(Q_k^{(l)} R_k^{(l)} \)是可学习的变换矩阵，k是SE3的表示维度。
2.3 解码器
2.3.1 结构预测头
用于预测三位坐标，二面角或距离矩阵。
\( \sigma^{(coord)} \)是激活函数，\( W_{c}^{(l)} \)是坐标回归权重。
2.3.2 序列设计头
3.2.1 自回归生成（交叉熵损失）
逐个生成氨基酸序列，条件于当前节点特征和已生成序列：
\( a{v}^{(t)} \)是时间t的氨基酸类型，\( W{s}^{(l)} W_{emb}^{(l)} \)为权重矩阵。
损失函数（序列恢复）
交叉熵损失对比预测序列与天然序列：
多任务学习
同时优化结构预测（如坐标误差）和序列设计（如恢复天然序列）任务，通过联合损失函数平衡不同目标。
三 实验
3.1 设备
Atlas 800T A2
3.2 组件版本
hdk：24.1.RC3
cann：8.0.RC3
python：3.10
torch：2.1.0
torch-npu：2.1.0.post8
3.3 环境准备
官方文档对依赖包没有版本要求，仅需保证python的版本高于等于3.0即可。
新建conda环境
conda create --name mlfold python=3.10
安装torch_npu和一些必要的依赖包
下载torch_npu
wget 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
注：该环境已经提前安装过了，所以不会提示安装成功数据集下载
3.4 下载原始数据集
wget https://files.ipd.uw.edu/pub/training_sets/pdb_2021aug02.tar.gz
注：原始pdb_2021aug02数据集压缩包大小为17GB，解压之后为92GB，须预留足够的空间。
下载之后使用tar -xzf pdb_2021aug02.tar.gz进行解压。
下载用于测试的数据集样本
如果没有足够的磁盘空间，可下载该数据集，压缩包为48MB，解压之后为255MB。
wget https://files.ipd.uw.edu/pub/training_sets/pdb_2021aug02_sample.tar.gz
3.5 推理
ProteinMPNN在example目录下面有很多样例，均可用来验证推理。
在training目录下也有test_inference.sh用于推理验证。
3.5.1 激活conda环境
conda activate mlfold
3.5.2 修改推理脚本
Slurm参数、用于输入的.PDB文件、权重路径、权重名称
3.5.3 修改protein_mpnn_run.py
3.5.4 执行推理
3.5.5 查看推理结果
3.6训练
3.6.1 修改training.py文件
注：training.py脚本内，使用的cuda api不需要手动修改。torch_npu会自动将cuda api更改为torch_npu的api。例：
3.6.2 修改sh脚本
3.6.3 source cann
3.6.4 验证npu环境是否可用
3.6.5 开始训练
在training目录下面执行bash submit_exp_020.sh
查看npu状态，显示0卡正在被进程占用
注：可以通过环境变量export ASCEND_RT_VISIBLE_DEVICES来指定用哪张卡训练。
3.6.6 训练完成
训练完成之后，训练日志和权重会存放在exp_020目录下。
