等变图神经网络在药物研发中大放异彩

在药物研发领域，AI需要对药物小分子和蛋白质处理带有几何特征的图。在这个图中，除了包含原子的一些内在特征以外，我们还需要考虑到每个原子在空间的三维坐标这一几何特征。不同于一般特征，这些几何特征往往都具备着一些对称性和等变性。等变图神经网络模型对这类等变对称性的特征可以很好的建模。

本文基于腾讯 AI Lab， 清华 AIR & 计算机系的综述:“Geometrically Equivariant Graph Neural Networks: A Survey”文章，引出等变图神经网络中的基本概念，并结合最近在药物研发领域中发表的文献，来分析讨论等变图神经网络的巨大应用价值。

具体的，先解释等变图神经网络的基本概念(包括群、等变的基本概念以及等变图神经网络的分类)，然后分享等变图神经网络在分子动力学模拟、分子生成与预测、蛋白质与蛋白质docking、蛋白质与小分子docking、蛋白质结构精调与评估、抗体序列与结构协同设计等的应用。

基于对对称性的建模，大量基于图神经网络的改进模型在近年来被提出。这一类模型，因为克服了传统图神经网络无法很好处理这类具有等变对称性质的特征的缺点，被统称为等变图神经网络。该综述文献对等变图神经网络的结构和相关任务进行了一个系统梳理。

设X和Y分别为输入和输出向量空间，这两个空间都被赋予一组变换G：G × X → X and G × Y → Y，函数φ : X → Y，如果当我们对输入应用任何变换时，输出也会通过相同的变换或者在某种可预测的行为下发生变化，则Y称为相对于G的等变。

等变(Equivariance)

函数φ : X → Y是G-等变的，如果X → Y计算的是在G上的任何变换

其中ρX是输入中的群表征和ρY是输出空间。具体来说，如果ρY是恒等式(identity)，则φ称为不变的(invariant)。

群(group)

群G是一组具有满足这些的二元运算"."属性的变换：这些属性存在一个标识元素，每个元素G存在一个逆，使得在associative composition下是闭集的。

给定群的定义，提供一些例子：

O(n)是由旋转和映射组成的n维正交群；

SO(n)是一个只包含旋转的特殊正交群；

E(n)是由旋转、映射和平移组成的n维欧几里得群；

SE(n)是一个由旋转和平移组成的特殊的欧几里得群；

Lie李群是其元素构成可微流形的群。实际上，上面所有的群都是李群的具体例子。

基于消息表示的类别不同，可以将等变图神经网络模型分为三大类：不可约表示(Irreducible Representation)，正则表达(Regular Representation)和标量化(Scalarization)。

基于不可约表示信息的模型

这类模型基于表示论中关于紧群的线性表示可以拆解为一系列的不可约表示的直积这一理论。从而在SE(3)群中构建满足等变性质的消息模型。例如，在TFN中

TFN 层利用 Clebsch-Gordan 系数的性质构造了一个对于任意旋转参数

，对于任意属于 SO(3)的旋转操作都等变。

有大量的工作基于 TFN 结构做了相应的扩展，例如加入 Attention 机制， 引入非线性的 Clebsch-Gordan 系数等。但是这类方法计算复杂度都较高，且不可约表示仅仅适用于特定的群。这约束了这类模型的表达能力。

基于正则表示信息的模型

另一类的工作尝试利用群的正则表示来构建群卷积操作。这里的代表性工作李卷积(LieConv)通过LiFiting操作将输入映射到群中的元素，然后利用PointConv完成群卷积的离散化计算。在我们的符号约定下，李卷积可以表示为：

其中

是

在群中的映射元素，log 将群元映射到对应的李代数，

是一个 MLP。通过这种构造，李卷积中对于h_i的更新实现了对于任意李群以及其离散子群的不变性。LieTransformer 基于此思想，引入了自注意力机制来进一步提高模型的性能。基于李群正则表示的模型在群的选取上更加灵活，但是由于要进行离散化和采样，需要在效率和性能之间做出权衡。同时，以上的更新只考虑了标量信息 h，但难以直接推广到对几何信息 x 的更新，除非综合哈密顿网络等工作中的更新方法。

标量化

除了基于群表示论的途径外，很多工作采用了一种基于标量化的建模等变性质的方法。这类标量化的方法将几何特征转化为一些不变的标量，然后利用MLP等网络结构来得到一个标量变化，最后将这个变化加回到原有的几何特征上从而得到等变性。

ECNN作为在标量化里面的一个重要工作，提出了一个非常灵活的框架：

其中，

是对几何特征的标量化，

函数为不同的 MLP，通过将几何信息和非几何信息消息进行关联

， EGNN 可以同时保证非几何特征和几何特征传播过程中的等变性。这个构造结合了物理知识，可以看成是对两个粒子的库伦力 / 重力的计算的建模。

在 EGNN 基础上，GMN 扩展了模型可以描述的几何特征维度，在建模坐标信息的同时也可以同时引入更多的几何信息（如速度、加速度、角速度等）并保证等变性。GemNet 则在 DimeNet 基础上通过这一通用的表示将一些更丰富的几何特征，例如二面角等，结合到消息传播的过程中。

01 等变图神经网络用于分子动力学模拟

2022月5月4日，哈佛大学Soris Kozinsky团队在nature communications上发表了题为"E(3)-equivariant graph neural networks for data-efficient and accurate interatomic protentials"的文章，开发了基于等变图神经网络模型NequIP ，该模型在分子动力学模拟上，显示出出色的预测准确性、对未见阶段的泛化能力以及非常高的样本效率。

该E(3)等变图神经网络方法：Neural Equivariant Interatomic Potentials (NequIP)，用于分子动力学模拟的从头计算中学习原子间势(interatomic potentials)。

现在大多数对称感知模型使用不变卷积并且只用于标量，NequIP使用E(3)等变卷积来处理几何张量的相互作用，从而产生更丰富的信息和更可信的原子环境表示。

NequIP在具有挑战性和多样化的分子和材料集上[包括小分子、不同相的水，无定形固体、固/气界面反应和锂超离子导体]实现了最优的准确性，同时表现出了最好的数据利用效率，训练数据减少多达三个数量，参数少于1000个甚至100个，该方法的高数据效率允许使用高阶量子化学理论水平作为参考构建准确的电位(accurate potentials)，并能够在长时间尺度上进行高保真分子动力学模拟。

本工作的贡献是为基于几何张量上的 E(3) 等变卷积的分子和材料引入了深度学习能量守恒的原子间势，从而产生了最先进的精度和出色的数据效率，并且可以从分子动力学模拟中以高保真度再现结构和动力学特性。

a.一组原子被解释为具有局部邻域的原子图。b.原子序数嵌入到 l=0特征中，通过一系列相互作用块进行细化，创建标量和高阶张量特征。一个输出块生成原子能量，这些原子能量被汇集后给出总预测能量。c.相互作用模块，包含卷积。d.卷积通过张量将径向函数R(r)和单位向量？？？的球偕投影与相邻特征的乘积相结合。

NequIP 网络架构，如上图所示，建立在原子嵌入的基础上，然后是一系列交互块，最后是输出块：

嵌入编码：在SchNet之后，初始特征是使用可训练的嵌入生成的，该嵌入仅对原子序数Zi(通过one-hot编码表示)进行操作，通过可训练的自交互层实现。

交互块：交互块编码相邻原子之间的交互；该块的核心是卷积函数。产生相同旋转和parity pair(lo, po)的不同张量积交互的特征，由线性原子级别自相互作用层混合，相互作用块由ResetNet更新。自相互作用权重是针对每个speices单独学习的。最后，混合特征由等变的基于SiLU的门非线性处理。

输出块：最终的卷积的 l=0 特征被传递到输出块，该输出块由一组两个原子级自相互作用组成。

每个原子，最后一层输出一个标量，被解释为原子势能。然后将这些相加得出系统的总预测势能。随后获的力作为预测的总势能的负梯度，从而确保能量守恒和旋转等变力。

虽然发现 NequIP 显示出出色的预测准确性、对未见阶段的泛化能力以及非常高的样本效率，但仍然存在的一个开放挑战是深度学习原子间势的可解释性。

经典原子间势中的能量贡献可以明确分配给个别类型的相互作用，例如pair-wise bonded terms或Coulomb或van der Waals non-bonded interactions。

仍然需要系统地探索将这些物理知识明确包含在深度学习原子间势的复杂功能形式中的潜在好处和最佳方法。另一方面，允许这种可解释性水平的经典力场函数形式的简单性严重限制了它们的准确性，在两种方法之间呈现出有趣的张力。

作者期望所提出的方法将使计算化学、物理学、生物学和材料科学的研究人员能够以更高的准确性和效率对复杂反应和相变进行分子动力学模拟。

代码路径：
https://github.com/mir-group/nequip

02 等变图神经网络用于分子生成、预测

分子中的原子会受到复杂化学相互作用的影响。对于分子数据，标量节点特征通常是原子数、节点之间的connectivity，这个connectivity要么由化学键提供，要么由基于距离阈值获取。等变图神经网络可以用于分子的生成、预测。具体任务包括Position predict、Chemical Property Prediction、Energy& Force prediction、Binding Affinity prediction、Classification。

7月18日腾讯 AI Lab，发表了题为“3D Equivariant Molecular Graph pretraining”文章，开发了基于能量表示的等变分子图神经网络预训练模型，在下游能量和力预测上均达到了最优的效果。

使用3D backbone模型，设计3D感知预训练任务，并针对下游任务进行评估。存在如下挑战：如何保持生物世界的对称性旋转/平移分子的三维构象不会改变其行为规律。在数学上，应该使用backbone模型E(3)等变和预训练损失E(3)不变，其中群E(3)收集旋转、映射和平移的转换。

作者提出了1).一个基于能量的表示模型，预测输入3D中每个原子的E(3)-等变分子，利用最近提出的等变GNN作为构建块；2). 为了预训练该模型，制定了一个受物理启发的节点级别的力预测任务，即进一步以等变的方式转化为三维位置去噪损失。最重要的是，利用提出了基于黎曼-高斯分布，开发了一种新的去噪方案，以确保预训练任务的E(3)等变；3). 在3D上额外设计了一个图级别的任务，能够提高最后的性能。

3D分子图预训练(3D-MGP)概述。它包括两个任务：等变节点级别的力预测和不变图级别噪声尺度的预测。

是和黎曼-高斯分布有关

在具有三维构象的大规模数据上预训练模型:GEOM-QM9，然后评估在两个流行的三维任务：MD17和QM9上的性能。大量实验表明，我们的模型优于最先进的2D方法，即使它们的输入增加了3D坐标以进行公平比较。

代码路径：https://anonymous.4open.science/r/3D-MGP-4931

2022年7月23日，MIT的Tess Smidt团队发表了题为"Equiformer:Equivariant Graph Attention Transformer for 3D Atomistic Graphs"的文章，作者将等变图神经网络结合Transformers的优势构建了Equiformer模型

与3D相关的归纳偏差(如平移不变性和旋转等变性)对于在3D原子图(如分子)上运行的图神经网络是必不可少的。受到Transformers在各个领域成功的启发，作者研究了将这些归纳偏差融入到Transformers中，提出了Equiformer，利用Transformers架构的优势并结合基于不可约表示(irreducible representations,irreps)的SE(3)/E(3)-等变图神经网络。

不可约表示特征在channel维度中编码等变信息，而不会使图结构复杂化。这种简化，使得能够将原始操作替换为等变对应操作来直接合并它们。此外，为了更好地使Transformers适应3D图，提出了新等变图注意力，该注意力同时包含了内容和几何信息。提高注意力的表达能力，将点积注意力替换为多层感知器注意力，并包括非线性消息传递。

两个量子性质预测数据集QM9和OC20上对Equiformer进行了基准测试。对于QM9，在使用相同数据分割训练的模型中，Equiformers在12个回归任务中的11个上取得了最佳结果。对于OC20，在使用IS2RE数据和可选的IS2RS数据进行训练的设置下，Equiformer改进了最先进的模型。

Equiformer的架构。利用原子的度和边的度的表征，用Transformers块编码3D图，Transformers块由等变图注意和前馈网络组成。”⊗” 表示乘法，”⊕” 表示加法和“DTP”表示深度张量积。圆内的P表示所有邻域的总和。灰色单元格表示中间不可约表示的特征

5月15日，清华人工智能研究所联合北大发表了题为"3DLinker:An E(3) Equivariant Variational Autoencoder for Molecular Linker Design"文章，开发了基于等变图变分自编码器，用于预测锚原子作为linker结合位点，并联合生成linker及其3D结构。

生成一个小的linker，以物理连接两个具有不同功能的独立分子。

计算挑战：1) 连接器的生成取决于两个给定的分子，这与之前从头开始生成完整分子相反；2) linker严重依赖于要连接的两个分子的锚原子，这是事先不知道的；3) 需要考虑分子的3D结构和方向以避免原子冲突，为此需要与E(3)群等价。

作者为了解决这些问题，提出了一个条件生成模型3Dlinker，它能够基于E(3)等变图变分自编码器预测锚原子并联合生成linker及其3D结构。3Dlinker模型在恢复分子图方面具有显著更高的速率，更重要的是，准确的预测了所有的3D坐标。

整体编码和解码过程的图示。对于编码，baselines：用MF-MP编码器获得节点级别表示。这些节点在片段中的表示将被丢弃，取而代之的是仅在片段图上单独计算的表示。对于解码，预测两个锚节点作为linker的结合位点。在连接之前，同时预测linker的节点类型。使用两个锚节点和linker的节点类型，可以顺序预测边喝坐标，如下图所示。

边和坐标的顺序预测图示。首先关注一个节点。然后，对关注节点和其他节点(包括人工停止的节点)之间的边进行采样。如果linker节点首先连接到现有图，则将预测其坐标。每次在预测之前，MF-MP都应用于从现有图中捕获信息。我们不断添加边，直到选择停止节点，然后同时更新现有图中所有linker节点的坐标。然后，我们重新关注一个新节点并重复。该过程将继续，直到linker中的所有节点都已关注。

代码路径：https://github.com/YinanHuang/3DLinker

03 等变图神经网络用于蛋白质与小分子、蛋白质与小分子docking

麻省理工学院的研究人员用SE(3)等变图神经网络，直接挑战传统docking方法过程中采样、打分、排序和微调的步骤，对蛋白质或小分子表征直接通过SE(3)等变图神经网络预测二元复合物的pose。

MIT作者第一篇工作"Independent SE(3)-Equivariant Models for End-to End Rigid Protein Docking"开发了EquiDock模型，直接将蛋白与蛋白对接速度提高500倍。

第二篇工作"EquiBind:Geometric Deep Learning for Drug Binding Structure Prediction"直接用EquiDock的等变图神经网络建模的思想，开发了EquiBind模型，直接将蛋白质与小分子对接的速度提高了1200倍，极大的提高了药物研发的速度。下面将具体揭开SE(3）等变图神经网络在docking应用中的神秘面纱。

蛋白质复合物的形成是生物学中一个核心问题，作者设计了一个新的成对SE(3)-等变图匹配网络来预测旋转和平移。

在数学上保证了一个基本原则：无论两个结构的初始位置和方向如何，预测的复合体总是相同的，模型EquiDock，使用关键点匹配和对齐，来逼近结合口袋并预测对接pose，这是通过优化传输Kabsch算法实现的。

因为不依赖大量候选采样、结构精调，中间省去了采样、打分、精调的过程，所以显著缩小了运行时间，并且优于现有对接软件。

预测了一个刚性对接，将配体放置在受体的结合位置。在数学上保证输出相同的复合物构象--直到SE(3)变换-独立于初始未结合位置，旋转，或者两个蛋白的角色(配体和受体互换)

IEGMN模型架构介绍

独立E(3)-等变图匹配网络(IEGMN)，Φ架构满足等式(4)称为独立E(3)-等变图匹配网络(IEGMN)。

IEGMN扩展了图匹配网络(GMN)和E(3)等变图神经网络。IEGMN输入图对G1=(V1, E1), G2=(V2, E2)执行节点坐标和特征嵌入更新，并使用节点间和节点内消息以及E(3)等变坐标更新。

IEGMNs将第l层节点特征编码

和节点坐标

转换为：

其中N(i)是节点i的邻居节点，ϕ x是实值参数函数；W是可学习矩阵。ϕ h , ϕe是输出向量R的参数函数(MLPs)。fj→i 和 fi是原始边和节点特征(从残基中SE(3)等变提取）.aj→i是一个基于注意力的系数，具有可训练的浅层神经网络 ψ q 和 ψ k。

ϕx , ϕh , ϕe , ψq , ψk为不同的IEGMN共享或不同层。几个IEGMN层输出表示为：

EQUIDOCK架构和损失函数的详细信息。a. IEGMN中的消息传递操作保证了两两独立SE(3)-等变。b. 使用额外的最优传输损失(OT, optimal transport loss)。预测与结合口袋位置对齐的每个蛋白质的关键点。c. 在预测对接位置后，计算配体上的MSE损失，以及防止两个蛋白交叉的损失。

a. 复合物RMSD(Complex-RMSD)分布(DIPS测试集)；b. Interface-RMSD分布；c. c-RMSD与I-RMSD的散点图。

EQUIDOCK复合物预测结果：优于基线。该方法不使用采样，排序，不依赖于手工特征，并且不执行结构微调，旨在直接预测SE(3)配体转换。

EQUIDOCK计算效率：推断时间比基线快80-500倍。这对旨在扫描广阔搜索空间的密集筛选应用程序尤其重要，例如药物发现。

代码地址：https://github.com/octavian-ganea/equidock_public

预测药物样分子如何与特定蛋白质靶标结合是药物发现的核心问题。一种极快的计算绑定方法将实现快速虚拟筛选或药物工程等关键应用。现有方法的计算成本很高，因为它们依赖于大量候选采样以及评分、排名和微调步骤。

作者用 EquiBind 挑战这一范式，这是一种 SE(3) 等变几何深度学习模型，对 i) 受体结合位置（盲对接）和 ii) 配体的结合姿势和方向进行直接预测。与传统和最近的基线相比，EquiBind 实现了1200倍的加速。

图中展示了EquiBind的模型架构。简单来说，该方法利用E(3)协变网络生成配体和受体上的结合关键点，而后在三维空间中对这些点进行匹配来确定结合模式。同时，该网络还能够对配体的三维结构进行变换实现小分子柔性对接。

从结果上看，EquiBind在各项指标上均显著优于其他模型。该优势主要体现在Centroid Distance：即EquiBind能够正确找到小分子的结合口袋。同时，EquiBind的运行时间也显著优于其他方法。

EquiBind是一种全新的基于深度学习的小分子对接方法，能够利用E(3)等变图神经网络快速确定蛋白-小分子的结合方式。该方法为小分子对接提供了一种全新的解决方案，有较好的发展空间。

尽管如此，该方法目前在评价上存在局限性。上述评价过程只涉及结合口袋未知的情况，但是大部分虚拟筛选项目所关注的口袋是确定的。

另外，本文所比较的baseline方法主要关注的是口袋已知的对接，在口袋未知的情况下一般表现较差，直接比较可能不公平。

后续工作可以评价EquiBind在口袋已知情况下的表现，并与相关baseline进行比较，从而进一步明确其适用范围。

代码路径：https://github.com/HannesStark/EquiBind

04 等变图神经网络用于蛋白质结构精调与评估

密苏里大学电气工程与计算机科学系特聘教授程健林团队，近日连发了两篇将等变图神经网络用于蛋白质结构评估的文章:“EGR:Equivariant Graph Refinement and Assement of 3D protein Complex Stuctures", "3D-equivariant graph neural networks for protein model quality assessment"。融合蛋白质结构精调过程和AlphaFold2的特征，使用等变图神经网络提高了蛋白质结果评估的效果。

作者开发了等变图精调模型(EGR)，一种E(3)等变图神经网络(GNN)，用于多任务结构精调和蛋白质复合物的评估。

文章的贡献：

(1) 提供了第一个将深度学习应用蛋白质复合物结构的全原子精调任务的示例。

(2) 提供了第一个将等变图消息传递应用于同时精调和评估蛋白质复合物结构的多任务设置的示例。

(3) 介绍了新的半监督EGR模型，展示了其在改善蛋白质链之间界面区域结构方面的有效使用，并估计了其在此类改进中的置信度。

代码：https://github.com/BioinfoMachineLearning/DeepRefine

EnQA整体架构示意图。输入模型的1D/2D特征首先被转换为3D等变图模块的隐藏节点和边特征。残基Ca原子的空间坐标也被用作额外的特征。节点和边网络模块迭代地更新图特征。最后，通过3D等变网络从更新的节点/边特征和空间坐标预测残基对的每个残基IDDT分数和距离误差。

作者利用了输入模型的几何结构特征和从 AlphaFold2 预测中提取的特征。该网络被开发为一个等变框架，节点和边特征通过节点和边级图网络。对不同结构模型数据集进行的计算实验证明 EnQA 实现了蛋白质质量评估的最先进性能。

更准确地说，在 CASP14 和最近的 CAMEO 蛋白质结构上，EnQA 在大多数评估指标上都优于所有其他方法，包括使用 AlphaFold2 预测作为评估模型的参考。

在评估高质量的 AlphaFold2 模型时，作者的方法比 AlphaFold2 的自我报告的 lDDT 得分表现更好。在所有测试数据集上，EnQA 的性能明显优于之前的 QA 方法，证明了使用 3D 等变架构和基于 AlphaFold2 的特征的价值。

此外，结果表明从结构模型中提取的输入特征与从 AlphaFold2 预测中提取的信息具有互补效应，特别是对于 EnQA 表现更好的模型。

05 等变图神经网络用于抗体设计

华盛顿大学联合清华、北大等团队发表了题为“Antigen-Specific Antibody Design and Optimization with diffusion-based Generative models“的文章，开发了基于扩散过程和等变神经网络联合模型，用于建模CDR序列和结构，可以明确地针对特定抗原结构并以原子分辨率生成抗体。

该模型是一把“瑞士军刀“，能够进行序列结构协同设计、给定骨架结构的序列设计和抗体优化。

作者额外构建了一个新的数据集，其中包括抗体样蛋白质，作为原始抗体数据集的补充。

作者的模型可以在通过生物物理能量函数和其他蛋白质设计指标测量的结合亲和力方面极具竞争力。

(a) 抗体-抗原复合物结构和CDR结构。(b) 氨基酸方向(用三角形表示)决定了它们的侧链方向，这是氨基酸相互作用的关键。(c)这项工作的任务是为给定的抗原结构和抗体框架设计CDR。

生成扩散过程的图示。在每一步，网络将当前的CDR状态作为输入，并将CDR的序列、位置、方向的分布参数化以用于下一步。最后，通过侧链打包算法构建全原子结构。

作者在抗体建模，各种任务的采样算法和新数据集管理方面的贡献：

(1) 提出了第一个深度学习模型，通过考虑抗原的3D结构来执行抗体序列结构设计。

(2) 在模型中，不仅设计蛋白质序列和坐标，还设计每个氨基酸的侧链方向(表示为SO(3)元素)。它是第一个可以实现原子分辨率抗体设计并且等变于旋转和平移的深度学习模型。

(3) 当前可以的复合结构相对稀缺，作者为社区策划了一个专注于抗体样蛋白洁结构的新数据集，以提高训练质量

(4) 模块可以应用于广泛的抗体设计任务，包括序列结构协同设计、固定骨架CDR设计和抗体优化。

该基于扩散的抗体设计生成模型，能够执行广泛的抗体设计任务，并且在所有这些任务上都可以实现高度竞争性的性能。

但是这项工作的也存在以下主要局限性: (1)它依赖于与目标抗原结合的抗体框架，(2)它集成了侧链填充算法，并且没有以端到端的方式生成全原子结构。

未来的工作包括研究如何在没有结合结构的情况下生成抗体，以及开发端到端的全原子生成模型。

参考文献

[1] Geometrically Equivariant Graph Neural Networks: A Survey

[2] E(3)-equivariant graph neural networks for data-efficient and accurate interatomic protentials

[3] 3D Equivariant Molecular Graph pretraining

[4] Equiformer:Equivariant Graph Attention Transformer for 3D Atomistic Graphs

[5] 3DLinker:An E(3) Equivariant Variational Autoencoder for Molecular Linker Design

[6] Independent SE(3)-Equivariant Models for End-to End Rigid Protein Docking

[7] EquiBind:Geometric Deep Learning for Drug Binding Structure Prediction

[8] EGR:Equivariant Graph Refinement and Assement of 3D protein Complex Stuctures

[9]3D-equivariant graph neural networks for protein model quality assessment

[10] Antigen-Specific Antibody Design and Optimization with diffusion-based Generative models