酷应用

GNN for Science: 腾讯AI Lab、清华共同发文综述等变图神经网络

百家作者：机器之心 2022-03-25 13:38:27

机器之心专栏

腾讯AI Lab，清华AIR&计算机系

一文了解等变图神经网络的结构和相关任务。

近年来，越来越多的人工智能方法在解决传统自然科学等问题上大放异彩，在一些重要的学科问题（例如蛋白质结构预测）上取得了令人瞩目的进展。在物理领域的研究中，非常多的物理问题都会涉及建模物体的的一些几何特征，例如空间位置，速度，加速度等。这种特征往往可以使用几何图这一形式来表示。不同于一般的图数据，几何图一个非常重要的特征是额外包含旋转，平移，翻转对称性。这些对称性往往反应了某些物理问题的本质。因此，最近以来，大量工作利用了几何图中的对称性，基于经典图神网络设计了很多具有等变性质的模型去解决对几何图建模问题。尽管在这一领域，等变图神经网络模型取得了长足的发展，但是还缺乏一个系统性的对这一领域的调研。为此，腾讯 AI Lab，清华 AIR & 计算机系在综述：《Geometrically Equivariant Graph Neural Networks: A Survey》中，对等变图神经网络的结构和相关任务进行了一个系统梳理。

综述论文链接：https://arxiv.org/abs/2202.07230

在物理和化学领域，很多问题需要去处理带有几何特征的图。例如，化学小分子和蛋白质都可以建模成一个有原子和其化学键关系组成的几何图。在这个图中，除了包含原子的一些内在特征以外，我们还需要考虑到每个原子在空间的三维坐标这一几何特征。而在物理学的多体问题中，每个粒子的几何特征则包括坐标，速度，旋转等。不同于一般特征，这些几何特征往往都具备着一些对称性和等变性。正因为如此，基于对对称性的建模，大量基于图神经网络的改进模型在近年来被提出。这一类模型，因为克服了传统图神经网络无法很好处理这类具有等变对称性质的特征的缺点，被统称为等变图神经网络。

在这篇综述里面，我们系统性的梳理了近年等变图神经网络的发展脉络，并且提供了一个简洁的视角帮助读者能够很快的理解这类网络的内涵。基于消息传播和聚合函数的不同，我们将现有的等变图神经网络分为三类。与此同时，我们还详尽阐释了当前的挑战和未来的可能方向。

等变图神经网络基础框架

在实际应用中，我们需要处理的图不仅包含拓扑连接和节点特征，同时也会包含一些几何特征。在使用图神经网络处理这些数据的时候，不同的特征需要满足不同的性质。例如，在预测分子的能量时，我们需要这个预测对于输入的几何特征是不变的，而在分子动力学应用中，我们则需要预测的结果和输入的几何特征是等变的。为了达到这样的目的，我们提出了一个等变图神经网络的通用框架：

在这个框架中，

代表输入图的几何特征，而h_i,h_j 代表非几何特征。

和

分别代表在边(x,j)上的几何和非几何的消息。

和

分别是针对几何和非几何消息的聚合函数。除此以外，针对非几何信息的消息函数

是对于输入

来说 G - 不变的。而针对几何信息的

消息函数则是对于输入

来说 G - 等变的。下图展示了这一通用框架的操作：

基于以上的通用框架，我们在下表总结了当前主流的等变图神经网络模型。与此同时，基于消息表示的类别不同，我们将现有的等变图神经网络模型分为三大类：不可约表示(Irreducible Representation)，正则表示(Regular Representation) 和标量化(Scalarization ).

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

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

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

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

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

基于正则表示信息的模型

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

其中

是

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

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

标量化

去基于群表示论的途径外，很多工作采用了一种基于标量化的建模等变性质的方法。这类标量化的方法先将几何特征转化为一些不变的标量，然后利用 MLP 等网络结构来得到一个标量变化，最后将这个变化加回到原有的几何特征上从而得到等变性。这种标量化的方法最早是由 SchNet 和 DimNet 提出，不过仅仅考虑了模型不变的部分。SphereNet 在之前工作的基础上进一步考虑了在消息传播网络上扭转角的变化。EGNN 作为在标量化里面一个重要的工作，提出了一个非常灵活的框架：

其中，

是对几何特征的标量化，

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

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

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

此外，还存在一类标量化的方法，其基于不变的标量和等变的向量的乘积仍然是等变的向量这一观察来构造等变的消息传播。例如， PaiNN 和 Equivariant Transformer 在不变的 SchNet 上通过径向基函数建模原子的距离将等变的性质扩展到 SchNet 上。

以下是对等变图神经网络模型的一个总结梳理：

等变图神经网络应用

因为可以更好的建模几何信息，等变图神经网络在从物理系统到化学物质的各种类型的现实世界几何数据中具有广泛的应用。这篇综述简单介绍其在物理系统，分子数据和点云数据上的应用。下表总结了现有的等变神经网络的应用方向和数据集：

对复杂物理系统的建模

长期以来，对复杂物理系统的动力学进行建模一直是一个具有挑战性的话题。在物理系统中，有像带电粒子这样的物体，它们通过基于物理定律的力进行相互作用产生运动轨迹。在 NRI 这篇工作中最早引入 n-body 模拟的问题。n-body 中系统包含多个带电粒子，这些粒子由相互之间的库伦力驱动。n-body 问题的目标是给定系统的初始条件（坐标，速度和电荷量）的情况下预测这些粒子的动力学轨迹。这一任务是 E(3)等变的。SE(3)-Transformer 和 EGNN 都展示了等变图神经网络在这个任务上的潜力。GMN 中进一步提出了一个更有挑战性的问题 -- 带约束的 n-body 问题，即如何在粒子之间有约束，例如连杆或者铰链的情况下对粒子运动轨迹做出有效的预测。除了微观方向的数据外，NRI 和 GMN 也采用了人体运动捕捉的宏观数据验证了模型的有效性。

对分子的建模

等变图神经网络另一个重要的应用方向是对分子数据的建模。在分子数据中，原子的相互作用是有一系列复杂的化学物理机制所决定的。对于典型的分子数据，原子非几何特征往往包含原子本身的一些特征，而几何特征则是原子的空间坐标，速度等。原子之间的边则由化学键或者根据实际距离做截断来构造。经典的在分子上的应用包括，分子预测和分子生成。

分子预测：具体来说，分子预测包含对分子的属性和结构的一些预测任务。在分子预测这一领域，包含以下经典数据集。在小分子方面：QM9 是一个包含 12 个量子特征预测任务的经典小分子数据集。M17 则是一个在 8 个小分子上得到的动态轨迹的数据集，其中还包含了对应状态的能量和相互力作用等信息。ISO17 则是一个类似的分子动态轨迹数据集，其包含了 129 个同分异构体的轨迹信息。The Open Catalyst 2020 (OC20) 则包含了催化剂和底物的催化过程的状态信息，其目标是给定初始状态预测目标结构和对应状态的能量。在大分子方面：MDAnalysis 是一个较为完备的关于蛋白质级别的分子动力学模拟的数据。Atom3D 是一个综合数据集，包含 8 个具有几何信息的分子预测任务，范围从小分子到 RNA 和蛋白质。

分子生成：在分子生成领域，等变图神经网络往往用于和分子构象相关的生成。ConfGF 和 DSGM 基于旋转 - 平移不变 GNN 来参数化打分函数并且构造了基于打分的构象生成模型。GeoDiff 则借助去噪扩散概率模型 (Denoising Diffusion Probabilistic Model) 并且基于具有等变保证的 GNN 来构造模型。等变流（Equivariant Flow）验证了基于等变核的标准化流（Normalizing Flow）的可行性。

对于点云的建模

点云是对象的一种表示格式，它通过一组分配有坐标的点来描述形状。在对点云建模这一领域，包含一些经典的数据和任务。ModelNet40 和 ScanObjectNN 是两个经典的点云数据集，其任务是做物品分类。因为在点云数据里面不存在显式的点和点的链接，在等变图神经网络建模点云的时候，往往会以一个距离 d 作为阈值来构造点和点之间的边。TFN 和 SE(3)-Transformer 都在点云数据取得了相对于传统方法具有竞争力的性能。

未来展望

在对现有等变图神经网络的方法和应用进行了系统性总结后。本综述也对这一领域未来潜在发展方向进行了一些讨论：

理论的完备性：不同于经典图神经网络，等变神经网络还缺乏一系列的关于表达能力和泛化性的理论分析框架。现有的一些工作主要集中在讨论消息传播机制中存在的通用表达。但是对于模型的整体性质依然不够清晰。如何构造一个完备的理论框架以指导模型的设计是未来十分有趣的方向。

大规模等变图神经网络，在前面已经提到，基于群表示理论的方法都具有计算复杂度太高这一缺点，这限制了等变图神经网络在大规模数据上的应用。尤其是在结合一些更复杂的结构，例注意力机制的时候，这一问题会更加严重。如何有效的精简现有模型并且加速计算，使得现有的等变图神经网络可以适用于大规模数据，是一个重要方向。

对多层次结构的建模：许多现实世界的系统都表现出复杂层次结构。例如，有机分子由多个官能团组成，蛋白质由氨基酸组成。通过利用这些结构，我们可以设计出对多粒度多层次结构的系统进行建模。与现有的只有单层结构的消息传递范式相比。这种层次结构的等变模型可能可以更好的刻画这种层次结构信息，提高模型的性能和泛化性。

新的应用和数据：现有的等变图神经网络模型大多数只是在规模和复杂性优先的系统上进行性能评估，例如模拟的 N-body 系统和小分子 MD 数据。未来，我们需要在一些更具有挑战性的任务上评估等变图神经网络的有效性，这些挑战包括：更多数量的对象，更复杂的交互，更多样化的约束等等。近期以来，在蛋白质建模上的工作是一些有益的尝试。但是由于数据收集的困难和数据质量的限制，尚未有一个全面的可以评估各种方法的数据集出现。在未来如何将等变神经网络扩展到更多，更复杂的领域，去解决现实问题是一个很有意义的方向。