ICLR 2022 | 三维分子图的球形信息传递

作者 | 庞超 编辑 | 李仲深

今天给大家介绍的是ICLR 2022 Poster的文章《Spherical Message Passing for 3D Molecular Graphs》。作者在此工作中考虑了三维分子图的表示学习，其中每个原子与三维的空间位置相关联。这是一个尚未得到充分探索的研究领域，目前还缺乏一个有效的信息传递框架。在这项工作中，作者在球坐标系(SCS)中进行了分析，以完整地识别三维图结构。基于此观察，作者提出了球形信息传递(SMP)作为一种新的和强大的三维分子学习方案。SMP显著降低了训练的复杂性，使其能够在大规模分子上有效地执行。此外，SMP能够区分几乎所有的分子结构，而未覆盖的案例在实际中可能并不存在。基于有意义的基于物理的三维信息表示，作者进一步提出了用于三维分子学习的SphereNet。实验结果表明，在SphereNet中使用有意义的三维信息可以显著提高预测任务的性能。结果还证明了SpherNet在可靠性、效率方面的优势。

一、介绍

在许多现实世界的研究中，像分子这样的物体被很自然地建模为图。随着深度学习的发展，图神经网络(GNNs)被人们提出，用于从图数据中学习知识。图神经网络中消息传递机制被广泛应用。在这项工作中，作者的目标是开发一种新的三维图的消息传递方法。一般来说，三维分子图包含笛卡尔坐标系中给出的每个原子的三维坐标以及图的结构。不同类型的相对三维信息可以从三维分子图中推导出来，它们在分子学习中很重要，如键长、键之间的角度等。

作者首先研究了三维分子的完整表示。这要求图结构必须由相对三维信息来唯一地定义。为此，作者在球坐标系(SCS)中进行了形式化分析，并表明每个原子在三维图中的相对位置是由三种几何形状唯一决定的，包括距离、角度和扭转角。然而，这种完整性需要基于边的2-hop信息，从而导致过高的计算复杂度。为了降低计算成本，作者提出了一种新的消息传递方案，称为球形消息传递(SMP)，用于快速和准确的三维分子学习。球形消息传递在表示三维分子方面是高效的和近似完整的。首先，作者设计了一种新的扭转角计算策略，它只考虑基于边的1-hop信息，从而大大降低了训练的复杂性。这使得球形消息传递能够推广到大尺度分子中。此外，作者还证明了球形消息传递可以区分几乎所有的三维图结构。未覆盖的情况被明确证明似乎很少在实践中出现。通过自然地使用相对的三维信息和一种新的扭转，球形消息传递可以生成对输入图的平移和旋转不变的预测。

二、分子的完整表示

相对三维信息可以是距离或角度，这些信息对输入分子的平移和旋转操作具有不变性。在球坐标系(SCS)中考虑这些信息是很自然的。作者首先研究了球坐标系中三维分子的结构鉴定。对于球坐标系中的任何点，其位置由三元组(d,θ,ϕ)指定，其中d,θ和ϕ分别表示径向距离、极角和方位角。当在球坐标系中建模三维分子图时，任何原子i都可以是局部球坐标系的原点，d、θ和ϕ则分别自然成为键长、键之间的角和扭转角。因此，原子i的每个相邻原子的相对位置可以由相应的元组(d,θ,ϕ)来指定。同样，可以确定每个原子在三维分子图中的相对位置，从而得到识别出的结构，输入图的平移和旋转操作不会改变这些信息。可以很容易地从笛卡尔坐标系转换到球坐标系，从而得到(d,θ,ϕ)。

如图1所示。作者利用过氧化氢的化学结构来展示d、θ和ϕ对分子结构的鉴定为什么是至关重要的。很明显，该结构仅由三个键长d1,d2,d3、两个键角θ1,θ2和扭转角ϕ来定义。注意，输入不能包含所有成对距离。这是因为原子的连通性通常是基于真实的化学键和截止距离。考虑所有的成对距离将导致严重的冗余，显著增加计算复杂度和预测误差。因此，在本例中，H-H键的长度不被认为是输入。显然，这两个O-H键可以围绕O-O键旋转，而不改变任何键长和键角。然而，在这种情况下，扭转角ϕ发生了变化，过氧化氢的结构也随之变化。扭转角的重要性也在相关的研究领域得到了证实。在这项工作中，作者提出球形消息传递来系统地考虑距离、角度和扭转角，以精确地表示学习三维分子图。

图1. 过氧化氢的化学结构

三、球形消息传递

3.1 消息传递方案

消息传递神经网络是神经网络中使用最广泛的体系结构之一。基于上文的完整性分析。作者建议在球形坐标系中执行消息传递，从而得到一种新的、有效的方案，称为球形消息传递。作者证明了现有方法中使用的消息传递方案，如SchNet和DimeNet，是球形消息传递的特殊情况。

首先作者正式定义了一个三维分子图，它通常被表示为一个四元组。是分子图G的全局特征向量。是原子特征的集合，其中每个都是原子i的特征向量。是边的集合，其中每个是特征向量，rk是接收原子的索引，sk是边k的发送原子的索引。是包含每个原子的三维空间信息的三维笛卡尔坐标集。此外，表示指向原子i的边集，表示原子i进入节点的索引。消息传递过程后的输出包括更新的全局特征向量、更新的原子特征和更新的边。

图2a中说明了球形消息传递的聚合方案。显然，原子rk的嵌入是通过聚合每个传入的消息ek来获得的。消息ek基于指向原子sk的传入消息集Esk进行更新。设q表示Esk中任何消息的发送方原子。可以定义一个局部球形坐标系，其中sk作为原点，消息ek的方向自然作为z轴。将sk的邻近原子o定义为参考原子。因此，参考平面由三个原子sk、rk和o形成。对于原子q，其位置由元组(d,θ,ϕ)唯一定义。d是它到原子sk的距离，θ指定其方向来更新消息ek。扭转角ϕ由定义的参考平面和由sk、rk和q跨越的平面形成。直观地说，球形消息传递作为三维图的高级消息传递架构，通过考虑所有距离、角度和扭转信息来指定任何相邻原子q的相对位置，从而对三维分子图进行更全面的表示。

图2. (a).球形消息传递的消息聚合方案。(b).在球形消息传递体系结构中计算扭转角的一个例子。

一般来说，原子sk可能有数个相邻的原子，表示为q1,...,qt。计算这些原子对应的键长和键角很容易。球形消息传递通过将这t个原子投影到垂直于ek并与sk相交的平面上来计算扭转角。然后在这个平面上，扭转角会在预定的方向上形成，如逆时针方向。通过这样做，任何原子自然都会在逆时针方向上成为其下一个原子的参考原子。这t个扭转角的和是2π。一个简化的情况如图2b所示。原子sk有三个相邻的原子q1,q2和q3；q3是q1的参考原子，它们形成ϕ1；q1是q2的参考原子，它们形成ϕ2；同样，q2是q3的参考原子，它们形成ϕ3。显然，ϕ1,ϕ2和ϕ3之和为2π。值得注意的是，将每个原子设计为预定义方向上的下一个原子的参考原子，这个设计是满足不变性的，因为参考原子是自然相对的。此外，此方法还计算了基于边的1-hop邻域内的扭转。即使一个扭转角涉及四个原子，此设计避免了扭转角的数量是指数级的，而是使它与相邻原子的数量相同。因此，它是高效的，且不会导致时间或内存问题。球形消息传递可以在球形坐标系中定义为

其中，φe,φv和φu分别是边、原子和整个图上的三个信息更新函数。ρe→v和ρv→u聚合不同类型的几何图形之间的信息。特别是在球形消息传递中，P中的3D信息被转换和合并以更新每个消息ek。因此，球形消息传递使用另一个位置聚合函数ρp→e来进行消息更新。

3.2 完整性与效率

上文描述的鉴定标准可以完全确定一个三维分子的结构，但它涉及到基于边的2-hop信息。因此，计算复杂度是O(nk3)，其中n为原子数，k表示每个中心原子的相邻原子的平均数目。不幸的是，这种设计很难推广到大尺度分子图。为此，作者提出了球形消息传递作为一种高效的、可扩展的方案来实现球形坐标系中的消息传递。球形消息传递只涉及基于边的1-hop信息，因此时间复杂度降低到O(nk2)。这使得球形消息传递能够应用于大分子。作者严格研究了球形消息传递的完备性，并表明它甚至可以区分复杂的几何性质，如手性，如图3a与3b所示。因为球形消息传递使用最后一个原子作为参考原子，在预定义的方向上，考虑相邻原子之间的相对顺序，而忽略绝对顺序。因此，球形消息传递不能区分如图3b和3c所示的两个分子。然而，这种情况在本质中可能不存在。这也在实验中证明，球形消息传递实现了与完整表示相当的性能，而后者存在巨大的时间复杂度和严重的内存问题。

图3：球形消息传递能够区分和不能区分的情况。sk的所有相邻节点都被投影到垂直于消息传递方向的平面上。假设所有的距离和角度都是固定的（否则分子可以被更容易区分）。因此，所示的所有角度都是扭转角，它们都是沿逆时针方向形成的。(a)和(b)是一对手性不同的分子，球形消息传递可以区分它们。球形消息传递不能区分(b)和(c)，但这种情况在本质中可能不存在。(b)中的∠q1skq2和(c)中的∠q1skq3通常不同，因为在距离和角度相同的情形下q2和q3通常是不同的原子。

3.3 与先前的消息传递方法的关系

具有相对信息的三维图的开发仍处于早期阶段。文献中有几种值得注意的方法，它们可以被看作是球形消息传递的特殊情况，因为它们捕获了部分三维位置信息。例如，SchNet和PhysNet考虑距离，而DimeNet对方向信息进行编码。在形式上，这些方法可以完全被球形消息传递的方程拟合。与之前的模型相比，球形消息传递提供了一个关于其完整性的严格证明，并明确描述了故障情况。重要的是，球形消息传递是基于三维分子图的识别分析而开发的。因此，它的目的是学习三维分子图的完整数据表示，而不是简单地包括额外的3D信息（如角度或扭转角）。

四、SphereNet

三元组(d,θ,ϕ)可以表示三维分子图中任何原子的相对位置。然而，它不能作为神经网络的直接输入，因为它缺乏有意义的表示。分子本质上是量子系统，因此分子表示需要遵循物理定律。一个重要的方面是选择适当的基函数，将三元组(d,θ,ϕ)转换为基于物理的表示。人们已经探索了几个基函数，包括MLP、高斯、正弦函数、球贝塞尔基函数和球谐函数。特别是球贝塞尔，被证明是编码距离的最佳基函数，而球形谐波是最适合编码角度的基函数。作者将最终表示记为Ψ(d,θ,ϕ)。这种表示的一种形式可以表示为，其中jl( )是l次的球贝塞尔函数，是l次m阶的球谐函数，c表示临界值，βln是l次贝塞尔函数的第n个根。l。L和N分别表示球谐函数和球贝塞尔函数的最高阶，它们是实验设置中的超参数。此外，还可以从Ψ(d,θ,ϕ)中推导出两个简化的表示Ψ(d)和Ψ(d,θ)。基于球形消息传递方案和物理表示，作者构建了用于分子学习的SphereNet。SphereNet可以产生既准确又有物理意义的数据表示。通过在球坐标中合并位置信息，SphereNet还生成对输入分子平移和旋转不变的预测。网络由一个输入模块、几个交互模块和一个输出模块组成。假设图中边k的消息ek是要更新的消息。具体地说，输入模块为边k生成初始消息，并且只使用距离表示Ψ(d)作为输入。每个交互模块都会更新边k的消息。它的输入包括所有相邻边的消息，以及所有三种表示，包括基于边k及其相邻边的Ψ(d,θ,ϕ)、Ψ(d,θ)和Ψ(d)。输出模块首先将k的距离表示和当前消息作为输入。边k的接收原子的特征向量是通过聚合指向它的所有消息来获得的，其中其他消息具有与ek类似的更新过程。

五、实验结果

表1. SphereNet与其他模型在IS2RE上在能量MAE和EwT真实能量的百分比方面的比较。最佳结果以粗体显示。

表2. SphereNet与其他模型在QM9数据集上在MAE和总体平均标准数据方面的比较。‘-’表示在参考论文中没有报道相应性质的结果。最好的结果用粗体显示，第二好的结果用下划线显示。

其他实验结果和模型的结构及参数详见原文。

六、总结

三维信息对分子很重要，但缺乏一个根据规则的信息传递框架来处理这个问题。作者首先研究了具有三维信息的分子的完整识别，并严格证明了一个完整的消息传递方案需要所有的距离、角度和扭转角信息。基于此类分析和效率的考虑，作者提出了球形信息传递作为一种统一而有效的处理三维分子图的方案。与原来的距离、角度和扭转角表示相比，球形信息传递显著降低了计算复杂度。重要的是，它也可以近似于分子的完整表示，因为失败的情况可能在自然界中不存在。基于球形信息传递和有意义的物理表示，作者提出了从三维分子图中学习的SphereNet。在各种类型的数据集上的实验结果表明，SphereNet具有显著的性能改进。此外，SphereNet可以推广到大尺度分子上快速执行而不降低性能。

原文链接

https://openreview.net/forum?id=givsRXsOt9r