重学ProteinMPNN原理1

ProteinMPNN

一、首先什么是MPNN for protein

机器学习中常见的蛋白数据表示

1.蛋白序列的独热编码
2.基于多序列比对的进化信息
3.二级结构信息
4.残基间距离
5.二面角扭转角度或其他结构物理信息
6.蛋白表面信息
7.原子坐标信息
8.蛋白质图（蛋白质的氨基酸残基被表示为图中的节点（vertices），而氨基酸之间的相互作用（如共价键、氢键、范德华力等）被表示为边（edges））

MPNN架构

训练过程
1.初始化
根据蛋白质结构中的位置和局部环境进行每一个氨基酸的初始化（为每个氨基酸残基分配一个特征向量）（也就是为什么输入要是完整结构而不能仅仅是backbone）
2.消息传递（分为聚合消息和节点更新两个步骤）
MPNN 通过消息传递机制来更新节点的嵌入。每个节点的嵌入会根据其邻居节点的信息进行更新。具体来说，每个节点会接收来自其邻居节点的消息，并通过聚合这些消息来更新自己的嵌入

消息传递中比较重要的两个参数选择
1.邻域大小：这篇文章中选择的邻域大小为30，也就是只选择空间上最近的30个氨基酸的信息进行传递
2.αC传递方向

3.自回归解码
对于节点级别的表示是不够的，需要获得氨基酸级别的表示
每个氨基酸的生成依赖于之前生成的氨基酸和当前节点的嵌入信息（其中，第一个氨基酸的生成是基于初始节点嵌入和结构信息进行的）

模型比较

与Rosetta固定骨架优化侧链协议在同一测试集上进行了比较：更快、序列有功能概率更高（Recovery）、对于复杂结构设计的序列可靠程度更高
注：Recovery衡量的是identity not similarity

二、ProteinMPNN

衡量了几种方法在蛋白核心（相邻8Cb距离近）和蛋白表面（相邻8Cb距离远）的序列恢复度，因为表面受到的物理约束更小，所以可能的氨基酸种类更多，恢复度较低。
proteinMPNN与2019的MPNN baseline模型相比，增加了以下三点改进：

①增加了除αC原子以外的几个原子之间的距离作为边的特征
②update edges：在消息传递过程中不止更新节点信息，还更新边的信息（两个氨基酸几个原子之间的欧几里得距离、一个氨基酸指向另一个的方向向量、两个氨基酸之间的二面角、氢键和范德华力等）
③在原始的N→C顺序解码的基础上增加了随机解码模式，使模型能学习到更广泛的序列生成模式（这样可以保证固定一些关键位置或者关键链的氨基酸，跳过他们生成其他位置的氨基酸）

评估指标除了recovery外还是添加了生成序列用AF2预测的精度

最后衡量了几种方法在蛋白核心（相邻8Cb距离近）和蛋白表面（相邻8Cb距离远）的序列恢复度，因为表面受到的物理约束更小，所以可能的氨基酸种类更多，恢复度较低。

PS:微调策略

找到train.py文件中的下面这段代码

model = ProteinMPNN(node_features=args.hidden_dim, 
                    edge_features=args.hidden_dim, 
                    hidden_dim=args.hidden_dim, 
                    num_encoder_layers=args.num_encoder_layers, 
                    num_decoder_layers=args.num_encoder_layers, 
                    k_neighbors=args.num_neighbors, 
                    dropout=args.dropout, 
                    augment_eps=args.backbone_noise)

后面加上冻结的层

# 冻结全部参数
for param in model.parameters():
    param.requires_grad = False

# 只解冻 decoder 层和输出层
for layer in model.decoder_layers:
    for param in layer.parameters():
        param.requires_grad = True

for param in model.W_out.parameters():
    param.requires_grad = True

写训练参数脚本，如

--previous_checkpoint model_weights/epoch_last.pt 
--path_for_training_data ./data/pdb_2021aug02 \
--path_for_outputs ./exp_test1 \
--num_epochs 100 \
--batch_size 8000 \
--hidden_dim 128 \
--num_encoder_layers 4 \
--num_decoder_layers 4 \
--dropout 0.1 \
--backbone_noise 0.2 \
--mixed_precision True

然后开始训练

参考资料
https://www.youtube.com/watch?v=6z4XmUAwdNA&t=80s
Robust deep learning–based protein sequence design using ProteinMPNN.2022
Generative models for graph-based protein design.2019