近年来,人工智能方法在解决传统自然科学问题上展现了巨大的潜力,特别是在一些重要学科(如蛋白质结构预测)中取得了显著进展。在物理研究中,许多问题涉及对物体几何特征的建模,如空间位置、速度和加速度等。这些特征通常可以用几何图来表示。不同于一般图数据,几何图的重要特性在于它们还包含旋转、平移和翻转对称性,这些对称性反映了某些物理问题的本质。因此,近年来,许多研究利用几何图的对称性,基于经典图神经网络设计出许多具备等变性质的模型来处理几何图的建模。尽管等变图神经网络模型在这一领域取得了显著进展,但仍缺乏系统性的研究综述。因此,腾讯 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结构进行了相应扩展,如引入注意力机制和非线性的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在点云数据上相较于传统方法取得了竞争力的性能。
未来展望,在对现有等变图神经网络的方法和应用进行系统性总结后,本综述对该领域未来潜在发展方向进行了讨论:理论的完备性。与经典图神经网络不同,等变神经网络仍缺乏一系列关于表达能力和泛化性的理论分析框架。现有一些工作主要集中在讨论消息传播机制中存在的通用表达,但对模型整体性质仍不够清晰。如何构建一个完备的理论框架以指导模型设计是未来的重要课题。
