蛋白质3D重建学习笔记：GNNRefine

一、代码Run

1. 运行

python /home/huangjiehui/Project/AIProtein/StructuralReconstruction/RefineGNN/ab_train.py --cdr_type 3 --train_path /home/huangjiehui/Project/AIProtein/StructuralReconstruction/RefineGNN/data/sabdab/hcdr3_cluster/train_data.jsonl --val_path /home/huangjiehui/Project/AIProtein/StructuralReconstruction/RefineGNN/data/sabdab/hcdr3_cluster/val_data.jsonl --test_path /home/huangjiehui/Project/AIProtein/StructuralReconstruction/RefineGNN/data/sabdab/hcdr3_cluster/test_data.jsonl

2. 评价指标

(1) RMSD值：**RMSD值 是两个分子之间（或者同一个分子的两个状态之间）的结构差异。分子对接方面则经常把结晶在蛋白质上的配体提取出来，拿对接软件再对接回去，然后比较对接前后的两个配体之间的RMSD值，这个值越小，说明对接软件的准确性更高。

(2)PPL值：困惑度（perplexity）的基本思想是：**给测试集的句子赋予较高概率值的语言模型较好,当语言模型训练完之后，测试集中的句子都是正常的句子，那么训练好的模型就是在测试集上的概率越高越好，**公式如下：

$PP(W)=P{\left(w_1{w}_2\dots w_N\right)}^{-\frac{1}{N}}=\sqrt[N]{\frac{1}{P(w_1{w}_2\dots w_N}}$

由公式可知，句子概率越大，语言模型越好，迷惑度越小。

(3)结果：我们报告重链中每个CDR的困惑度（PPL）和根均方差（RMSD）。右：抗原结合抗体设计任务的结果。我们报告所有方法的氨基酸回收（AAR）。10个epoch的结果：Test PPL = 8.625, Test RMSD = 2.503 (Training:3896, Validation:403, Test:437)

2.1 RMSD专栏

在动力学模拟分析中,很基本的一项分析就是RMSD值的求算。RMSD值即均方根偏差（Root Mean Square Deviation）。在统计学上，这个量就相当于标准差,反映的是数据偏离平均值的程度。在蛋白质结构解析，模建，结构联配(structure alignment）以及分子动力学模拟中，RMSD值是非常常用的一项参数,用于衡量原子偏离比对位置的程度。

在本例中，RMSD值反映的是泛素各个部分原子偏离平均位置的程度，也就是各原子运动幅度的大小。RMSD值越大，说明该原子的运动的空间范围越大,原子的空间位阻也就越小。原子α在某一段时间Nt内的RMSD值求算公式为:

$RMS{D}_{\alpha }\left(t_j\right)=\sqrt{\frac{{\sum }_{\alpha =1}^{N_{\alpha }}{\left(r_{\alpha }\left(t_j\right)-⟨r_{\alpha }⟩\right)}^2}{N_{\alpha }}}$

上式中，Na 是參与比较的原子数目，$r_{\alpha }\left(t_j\right)$ 是原子 $\alpha$ 在 $t$ 时刻的位置， $⟨r_{\alpha }⟩$ 是在整个时间段内原子 $\alpha$ 的平均位置。定义为：

$⟨r_{\alpha }⟩=\frac{1}{N_t}{\sum }_{j=1}^{N_t}{r}_{\alpha }\left(t_j\right)$

二、代码理解

1. 摘要

抗体是多功能蛋白质，与病毒等病原体结合并刺激适应性免疫系统。抗体结合的特异性由这些Y形蛋白质尖端的互补性决定区（CDR）决定。本文提出了一种生成模型，用于自动设计具有增强结合特异性或中和能力的抗体的CDR。之前的生成方法将蛋白质设计表述为结构条件序列生成任务，假设先验地给出了所需的3D结构。相比之下，我们建议将CDR的序列和3D结构作为图表共同设计。我们的模型在迭代细化其预测的全局结构的同时，自动回归地解开序列。推断的结构反过来指导后续的残留物选择。为了提高效率，我们以粗粒度的方式模拟CDR内部和外部残留物之间的条件依赖性。我们的方法在测试集上实现了卓越的对数可能性，并在设计能够中和SARS-CoV-2病毒的抗体方面优于之前的基线1

2. 引言

2.1 背景

单克隆抗体越来越多地被用作针对SARS-CoV-2等多种病原体的治疗方法（Pinto等人，2020年）。由于这些Y形蛋白质的结合特异性在很大程度上取决于它们的互补性确定区域（CDR），计算抗体设计的主要目标是自动创建具有所需特性的CDR子序列。由于2060多个可能的CDR序列的组合搜索空间和满足结合亲和力、稳定性和合成性所需约束的求解空间很小，这个问题尤其具有挑战性（Raybould等人，2019年）。

CDR生成中有三个关键的建模问题。首先是如何模拟序列与其底层3D结构之间的关系。生成没有相应结构的序列（Alley等人，2019年；Shin等人，2021年）可能会导致低于最佳性能（Ingraham等人，2019年），而从预定义的3D结构生成（Ingraham等人，2019年）不适合抗体，因为所需的结构很少先验地知道（Fischman和Ofran，2018年）。因此，开发共同设计序列和结构的模型至关重要。第二个问题是，如何根据序列的剩余部分（上下文）对CDR的条件分布进行建模。基于注意力的方法仅在序列级别模拟条件依赖性，但CDR与其上下文之间的结构相互作用对生成至关重要。最后一个问题与模型针对各种属性进行优化的能力有关。传统的基于物理的方法（Lapidoth等人，2015年；Adolf-Bryfogle等人，2018年）专注于结合能最小化，但在实践中，我们的目标可能比结合能更复杂（Liu等人，2020年）

2.2 本文亮点

本文将序列结构对表示为图，并将共同设计任务表述为图生成问题。图形表示允许我们在序列和结构级别上模拟CDR与其上下文之间的条件依赖关系。抗体图的生成带来了独特的挑战，因为当插入新节点时，全局结构预计将发生变化。以前的自回归模型（You等人，2018年；Gebauer等人，2019年）不能修改生成的结构，因为它们是在教师强迫下训练的。因此，在前面的步骤中犯的错误可能会导致后续生成步骤中的一系列错误。为了解决这些问题，我们提出了一种新的架构，将氨基酸节点的生成与3D结构的预测交织在一起。结构生成基于全局图的迭代细化，而不是用教师强迫对部分图的顺序展开。由于上下文序列很长，我们通过将节点分组为块来进一步引入粗图表示。我们在更粗糙的级别上应用图卷积，以有效地将上下文信息传播到CDR残基。在对具有已知结构的抗体进行预训练后，我们使用预定义的属性预测器对其进行微调，以生成具有特定性质的抗体

我们评估了我们三代任务的方法，从语言建模到SARS-CoV-2中和优化和抗原结合抗体设计。我们的方法与标准序列模型（Saka等人，2021年；Akbar等人，2021年）和针对抗体量身定制的最新图形生成方法（You等人，2018年）进行了比较。我们的方法不仅在测试序列上实现了较低的困惑，而且在属性引导抗体设计任务中的表现也优于之前的基线

3. 相关工作

3.1 抗体/蛋白质设计

目前的计算抗体设计方法大致分为两类。第一类基于能量函数优化（Pantazes & Maranas，2010年；Li等人，2014年；Lapidoth等人，2015年；Adolf-Bryfogle等人，2018年），该优化使用蒙特卡洛模拟迭代修改序列及其结构，直到达到局部能量最小值。蛋白质设计中使用了类似的方法（Leaver-Fay等人，2011年；Tischer等人，2020年）。然而，这些基于物理的方法在计算上很昂贵（Ingraham等人，2019年），我们期望的目标可能比低结合能复杂得多（Liu等人，2020年）。

第二类基于生成模型。对于抗体，它们主要是基于序列的（Alley等人，2019年；Shin等人，2021年；Saka等人，2021年；Akbar等人，2021年）。对于蛋白质，O’Connell等人。（2018年）；Ingraham等人。（2019）；Strokach等人。（2020）；Karimi等人。（2020年）；Cao等人（2021年）进一步开发了基于骨干结构或蛋白质褶皱的模型。我们的模型还寻求整合用于抗体生成的3D结构信息。由于新病原体通常不知道最好的CDR结构，我们共同设计了特定属性的序列和结构

3.2 图形的生成模型

我们的工作与图形生成的自回归模型有关（You等人，2018年；Li等人，2018年；Liu等人，2018年；Liao等人，2019年；Jin等人，2020年a）。特别是，Gebauer等人（2019）开发了用于分子图和构象共同设计的G-SchNet。与我们的方法不同，它们按顺序生成边缘，当新节点到达时，无法修改之前生成的子图。虽然Graphite（Grover等人，2019年）也使用迭代细化来预测图的邻接矩阵，但它假设给出了所有节点标签，并且只预测边缘。相比之下，我们的工作将自回归模型与迭代细化相结合，生成带有节点和边缘标签的完整图，包括节点标签和坐标

3.3 3D结构预测

我们的方法与蛋白质折叠密切相关（Ingraham等人，2018年；Yang等人，2020年a；Baek等人，2021年；Jumper等人，2021年）。对AlphaFold等最先进模型的输入需要完整的蛋白质序列、多序列对齐（MSA）及其模板功能。这些模型不直接适用，因为我们需要预测不完整序列的结构，并且没有提前指定MSA

我们的迭代细化模型还与分子构象预测的分数匹配方法（Shi等人，2021年）和点云基于扩散的方法（Luo & Hu，2021年）有关。这些算法还迭代细化了预测的3D结构，但仅适用于完整的分子或点云。相比之下，我们的方法学习预测不完整图的3D结构，并将3D结构细化与图形生成交织在一起

4. 具体路线

4.1 总览

**抗体序列和结构共同设计：**抗体的作用是与抗原（如病毒）结合，将其呈现给免疫系统，并刺激免疫反应。一种被称为中和抗体的抗体子集不仅与抗原结合，还可以抑制其活性。抗体由重链和轻链组成，每个链都由一个可变结构域（VH/VL）和一些常数结构域组成。可变域进一步分为一个框架区域和三个确定互补的区域（CDR）。重链上的三个CDR被标记为CDR-H1、CDR-H2、CDR-H3，每个CDR都占据一个连续的子序列（图1）。作为抗体最可变的部分，CDR是结合和中和的主要决定因素（Abbas等人，2014年）。
![在这里插入图片描述](https://img-blog.csdnimg.cn/a5812009939d4e1a5e5ef1de4460b5f.png#pic_center

图1. 抗体的示意图结构（从维基百科修改的图）

关注Shin等人。（2021年）；Akbar等人（2021年），我们制定抗体设计作为CDR生成任务，以框架区域为条件。具体来说，我们将抗体表示为图表，该图编码其序列和3D结构。我们提出了一种名为RefineGNN的新图形生成方法，并针对固定框架区域将其扩展到处理条件生成。最后，我们描述了如何将RefineGNN应用于属性引导优化，以设计具有更好中和性能的新抗体。为了简单起见，我们专注于生成重链CDR，尽管我们的方法可以很容易地扩展到模型轻链CDR。

符号:
抗体VH结构域表示为氨基酸序列s = s1s2 · ·· sn。序列中的每个令牌si被称为残基，其值可以是20种氨基酸之一，也可以是特殊的代币hMASKi，这意味着其氨基酸类型未知，需要预测。VH序列折叠成3D结构，每个残基si都标有三个主干坐标：xi；α表示其α碳原子，xi；c代表其碳原子，xi;n用于其氮原子。

4.2 图形表示

我们将抗体（VH）表示为图G(s) = (V; E)，节点特征V = {v1; · · · }，边缘特征$sqrtabc$ 。每个节点的特征vi编码三个二面角（φi; i; !I）与残留物的三个骨干坐标有关i。对于每个残留物i，我们计算一个方向矩阵Oi，表示其局部坐标系（Ingraham等人，2019年）（在附录中定义）。这使我们能够计算描述两个残基i和j之间空间关系的边缘特征：$\mathcal{E}={\left\{{\mathbf{e}}_{ij}\right\}}_{i\ne j}$ ，每个节点$v_i$会对三个二面角$\left({\varphi }_i,{\psi }_i,{\omega }_i\right)$进行编码，他们与残留物i的三个骨干坐标 $i$ 有关。对于每个残留物i，我们计算一个方向矩阵Oi，表示其局部坐标系（Ingraham等人，2019年）（在附录中定义）。这使我们能够计算描述两个残基i和j之间空间关系的边缘特征：

${\mathbf{e}}_{ij}=\left(E_{\mathrm{pos}}(i-j),\mathrm{RBF}\left(\Vert {\mathbf{x}}_{i,\alpha }-{\mathbf{x}}_{j,\alpha }\Vert \right),{\mathbf{O}}_i^{\top }\frac{{\mathbf{x}}_{j,\alpha }-{\mathbf{x}}_{i,\alpha }}{\Vert {\mathbf{x}}_{i,\alpha }-{\mathbf{x}}_{j,\alpha }\Vert },\mathbf{q}\left({\mathbf{O}}_i^{\top }{\mathbf{O}}_j\right)\right)$

边缘特征eij包含四个部分。位置编码Epos（i − j）编码抗体序列中两个残基之间的相对距离。第二项RBF（·）是一种提升为径向基的距离编码，Eij中的第三个项是方向编码，对应于残留物i的局部框架中xj的相对方向。最后一个术语q（Oi>Oj）是空间旋转矩阵Oi>Oj的四元数表示q（·）的方向编码。

4.3 迭代细化图神经网络

（RefineGNN）：我们建议通过迭代细化过程生成抗体图。设G（0）是真实抗体图的初步猜测。每个残基被初始化为特殊令牌MASKi，每个边缘（i；j）初始化为距离为3ji-jj，因为连续残基之间的平均距离约为3。方向和方向特征设置为零。在每一代步骤t中，模型都学习修改当前的抗体图G(t)，并预测下一个残基t + 1的标签。具体来说，它首先使用参数θ的消息传递网络（MPN）对G（t）进行编码：

$\left\{{\mathbf{h}}_1^{(t)},\cdots ,{\mathbf{h}}_n^{(t)}\right\}={\mathrm{MPN}}_{\theta }\left({\mathcal{G}}^{(t)}\right)$

其中h(t) i是当前图G(t)下残留物i的习得表示。我们的MPN由具有以下架构的L消息传递层组成

${\mathbf{h}}_i^{(t,l+1)}=\mathrm{LayerNorm}\left({\sum }_j\mathrm{FFN}\left({\mathbf{h}}_i^{(t,l)},{\mathbf{h}}_j^{(t,l)},E\left({\mathbf{s}}_j\right),{\mathbf{e}}_{i,j}\right)\right),0\le l\le L-1$

其中h(t;0) i = vi，h(it) = h(it;L)。FFN是一个具有ReLU激活功能的双层前馈网络（FFN）。E（sj）是氨基酸类型sj的学习嵌入。根据学习的残留物表示，我们预测下一个残基t + 1的氨基酸类型（图2A）

图2：（A-C）RefineGNN的一代步骤。每个圆圈代表CDR残留物，每个正方形代表粗上下文序列中的残留块。（D）序列粗化。

$p_{t+1}=\mathrm{softmax}\left({\mathbf{W}}_a{\mathbf{h}}_{t+1}^{(t)}\right)$

这个预测给了我们一个新的图G(t+0.5)，其边缘与G(t)相同，但t+1的节点标签发生了变化（图2B）。接下来，我们需要更新结构，以适应新的残留物t+1。为此，我们用另一个具有不同参数θ~的MPN对图G（t+0.5）进行编码，并预测所有残留物的坐标。

$\begin{array}{rl}\left\{{\mathbf{h}}_1^{(t+0.5)},\cdots ,{\mathbf{h}}_n^{(t+0.5)}\right\}& ={\mathrm{MPN}}_{\tilde{\theta }}\left({\mathcal{G}}^{(t+0.5)}\right)\\ {\mathbf{x}}_{i,e}^{(t+1)}& ={\mathbf{W}}_x^e{\mathbf{h}}_i^{(t+0.5)},1\le i\le n,e\in \{\alpha ,c,n\}\end{array}$

新坐标x（t+1）i为下一次迭代定义了一个新的抗体图G（t+1）（图2C）。我们明确地实现了每个残留物的坐标，因为我们需要计算G（t+1）的空间边缘特征。结构预测（坐标xi）和序列预测（氨基酸类型pt+1）由两种不同的MPN进行，即结构网络θ~和序列网络θ。这种解体使两个网络能够专注于两个不同的任务。

培训。在培训期间，我们只应用教师强迫进行离散氨基酸类型预测。具体来说，在每代步骤t中，残基1到t被设置为其真实的氨基酸类型s1; · · · ; st，而所有未来的残基t + 1; · · · ; n设置为填充令牌。相比之下，连续结构预测是在没有老师强迫的情况下进行的。在每次迭代中，该模型都完善了上一步预测的整个结构，并根据预测坐标构造了所有残留物的新K近邻图G（t+1）$\left\{{\mathbf{x}}_{i,e}^{(t+1)}\mid 1\le i\le n,e\in \{\alpha ,c,n\}\right\}$

损失函数。我们的模型保持旋转和平移不变，因为损失函数是通过成对距离和角度而不是坐标计算的。抗体结构预测的损失函数由三部分组成。

• 距离损失：对于每个残留物对i, j，我们计算其预测的α碳原子 ${\mathbf{x}}_{i,\alpha }^{(t)},{\mathbf{x}}_{j,\alpha }^{(t)}$ 之间的成对距离。我们将距离损失定义为预测和真实成对距离之间的Huber损失:

  ${\mathcal{L}}_d^{(t)}={\sum }_{i,j}{\ell }_{\text{huber }}\left({\Vert {\mathbf{x}}_{i,\alpha }^{(t)}-{\mathbf{x}}_{j,\alpha }^{(t)}\Vert }^2,{\Vert {\mathbf{x}}_{i,\alpha }-{\mathbf{x}}_{j,\alpha }\Vert }^2\right)$

  距离为平方，以避免导致数值不稳定的平方根运算

• 二面体角损失：对于每个残留物，我们计算其二面角φ(it); i(t); !I(t) 基于预测的原子坐标 x( i;α t) ; x(i;c t); x(i;n t) 和 x(i+1 t) ;α; x(i+1 t) ;c; x(i+1 t) ;n。我们将二面角损失定义为预测和真实二面角之间的平均平方误差

  ${\mathcal{L}}_a^{(t)}={\sum }_i{\sum }_{a\in \{\varphi ,\psi ,\omega \}}{\left(\cos a_i^{(t)}-\cos a_i\right)}^2+{\left(\sin a_i^{(t)}-\sin a_i\right)}^2$

• Cα角度损失：我们计算两个向量x(i−t)1;α − x(i;α t)和x(i;α t) − x(i+1 t) ;α之间的双面角γi(t)，以及由x(i−t)2;α; x(i−t)1;α; x(i;α t) ; x(i+1 t) ;α定义的两个平面之间的二面角βi(t)。

  ${\mathcal{L}}_c^{(t)}={\sum }_i{\left(\cos {\gamma }_i^{(t)}-\cos {\gamma }_i\right)}^2+{\left(\cos {\beta }_i^{(t)}-\cos {\beta }_i\right)}^2$

总之，整体图形生成损失被定义为L = Lseq + Lstruct，其中 ${\mathcal{L}}_{\text{struct }}={\sum }_t{\mathcal{L}}_d^{(t)}+{\mathcal{L}}_a^{(t)}+{\mathcal{L}}_c^{(t)}{\mathcal{L}}_{\text{seq }}={\sum }_t{\mathcal{L}}_{ce}\left({\mathbf{p}}_t,s_t\right)$，序列预测损失Lseq是预测和真实残留类型之间的交叉熵Lce。

4.4 给定框架区域的条件生成

到目前为止描述的模型架构是为无条件生成而设计的——它不受任何约束地生成整个抗体图。在实践中，我们通常固定抗体的框架区域，并仅设计CDR序列。因此，我们需要扩展模型体系结构，以学习条件分布P（s0jsr），其中sr = sr+1 · · · sn是CDR sl之外的残留物 · · · sr。

(1)通过注意力调节: RefineGNN的一个简单扩展是使用循环神经网络对非CDR序列进行编码，并通过注意力层将信息传播到CDR。具体来说，我们首先将sr串联成上下文序列s~ = sr，其中⊕表示字符串串联，hMASKi重复n次。然后，我们通过门控循环单元（GRU）编码此上下文序列（Cho等人，2014年），并将结构和序列预测步骤（方程4和6）修改为

(2)多分辨率建模: 仅凭基于注意力的方法是不够的，因为它没有模拟上下文序列的结构，从而忽略了其残留物如何与CDR的结构相互作用。虽然在测试时新抗体无法获得这些信息，但我们可以学习使用具有已知结构的训练集中的抗体来预测这种相互作用；

一个天真的解决方案是迭代地细化整个抗体结构（超过100种残基），同时产生CDR残基。这种方法在计算上很昂贵，因为我们需要重新计算每代步骤中所有残留物的MPN编码。重要的是，我们无法从一开始就预测上下文残留坐标并修复它们，因为当CDR残留坐标在每个生成步骤中更新时，它们需要进行相应调整。

为了提高计算效率，我们提出了一个粗粒度模型，该模型通过将其集群到残余块中来减少上下文序列长度。具体来说，我们通过将每个b上下文残差聚类到块中来构建一个粗上下文序列bl;r(s)（图2D）。新序列bl;r(s)在残余块上定义了一个粗图G(bl;r(s))，其边缘是根据块坐标定义的。每个块xbi;e的坐标被定义为块内残留物的平均坐标。每个块E（bi）的嵌入是其残留嵌入的平均值。

$E\left({\mathbf{b}}_i\right)={\sum }_{{\mathbf{s}}_j\in {\mathbf{b}}_i}E\left({\mathbf{s}}_j\right)/b,{\mathbf{x}}_{{\mathbf{b}}_i,e}={\sum }_{{\mathbf{s}}_j\in {\mathbf{b}}_i}{\mathbf{x}}_{j,e}/b,e\in \{\alpha ,c,n\}$

现在，我们可以应用RefineGNN生成CDR残基，同时通过预测所有块的坐标迭代细化全局图G(bl;r(s))。唯一的变化是结构预测损失是在块坐标xbi;e上定义的。最后，我们结合了注意力机制和粗粒度建模，以保留细粒度和粗粒度信息。算法1说明了这种条件RefineGNN的解码过程。

4.5 属性引导的序列优化

我们的最终目标是产生具有所需特性的新抗体，例如中和特定病毒。这项任务可以表述为优化问题。设Y为中和的二进制指示变量。我们的目标是学习一个条件生成模型PΘ(s0jbl;r(s))，该模型最大限度地提高一组训练抗体D的中和概率，即

${\sum }_{\mathbf{s}\in \mathcal{D}}\log P\left(Y=1\mid {\mathbf{b}}_{l,r}(s)\right)={\sum }_{\mathbf{s}\in \mathcal{D}}\log {\sum }_{{\mathbf{s}}^{\prime }}f\left(s^{\prime }\right)P_{\Theta }\left(s^{\prime }\mid {\mathbf{b}}_{l,r}(\mathbf{s})\right)$

其中f(s0)是P的预测器（Y = 1js0）。假设给出了f，这个问题可以通过迭代目标增强（ITA）来解决（Yang等人，2020b）。在ITA优化开始之前，我们首先在一组真实抗体结构上预训练我们的模型，以学习CDR序列和结构的先验分布。在每个ITA微调步骤中，我们首先从D中随机抽取序列s，D是一组需要重新设计的抗体。接下来，我们根据上下文bl;r(s)生成M新序列。如果预测生成的序列s0 i被预测为中和，则将其添加到我们的训练集Q中。最初，训练集Q包含已知中和的抗体（Y = 1）。最后，我们从Q中取样一批中和抗体，并通过最小化其序列预测损失Lseq（Eq.10））。ITA微调阶段排除了结构预测损失Lstruct，因为生成序列的结构未知。