当前位置: 首页 > 科技观察

GNNforScience:腾讯AILab、清华大学等变量图神经网络联合发表综述

时间:2023-03-12 08:48:32 科技观察

近年来,越来越多的人工智能方法在解决传统的自然科学问题,以及一些重要的学科问题上大放异彩(对于例如,在蛋白质结构预测方面取得了令人瞩目的进展)。在物理学领域的研究中,很多物理问题都涉及到对物体的一些几何特征进行建模,比如空间位置、速度、加速度等,这些特征往往可以用几何图形的形式来表示。与一般的图形数据不同,几何图形的一个非常重要的特征是它们还包含旋转、平移和翻转对称性。这些对称性通常反映了某些物理问题的性质。因此,最近有大量工作利用几何图的对称性,在经典图神网络的基础上设计了许多具有等变特性的模型来解决几何图的建模问题。尽管等变图神经网络模型在该领域取得了长足的进步,但该领域仍缺乏系统的研究。为此,TencentAILab,TsinghuaAIR&ComputerDepartment对equivariantgraphneuralnetwork的结构和相关任务在综述中进行了系统回顾:《Geometrically Equivariant Graph Neural Networks: A Survey》。综述论文链接:https://arxiv.org/abs/2202.07230在物理和化学领域,很多问题都需要处理具有几何特征的图。例如,小的化学分子和蛋白质可以建模为由原子及其化学键关系组成的几何图。在这个图中,除了包括原子的一些本征特性外,我们还需要考虑每个原子在空间中的三维坐标的几何特性。在物理学的多体问题中,每个粒子的几何特征包括坐标、速度、旋转等。与一般特征不同,这些几何特征往往具有一定的对称性和等变性。正因为如此,基于对称性的建模,近年来提出了大量基于图神经网络的改进模型。这类模型由于克服了传统图神经网络无法处理此类具有等变对称特性的特征的缺点,所以统称为等变图神经网络。在这篇综述中,我们系统地梳理了近年来等变图神经网络的发展,并提供了一个简洁的视角,帮助读者快速理解此类网络的内涵。基于消息传播和聚合函数的差异,我们将现有的等变图神经网络分为三类。同时,我们阐述了当前的挑战和未来可能的方向。等变图神经网络基本框架在实际应用中,我们需要处理的图不仅包含拓扑连接和节点特征,还包含一些几何特征。在使用图神经网络处理这些数据时,不同的特征需要满足不同的性质。例如,在预测分子的能量时,我们需要预测对输入几何结构不变,而在分子动力学应用中,我们需要预测结果对输入几何结构是等变的。为了实现这个目标,我们提出了一个等变图神经网络的通用框架:在这个框架中,表示输入图的几何特征,而h_i,h_j表示非几何特征。并分别表示边(x,j)上的几何和非几何消息。和分别是几何和非几何消息的聚合函数。此外,非几何信息的消息函数对于输入是G不变的。几何信息的消息函数关于输入是G等变的。下图展示了这个通用框架的运行:基于上述通用框架,我们将目前主流的等变图神经网络模型总结如下表。同时,根据消息表示的不同类别,我们将现有的等变图神经网络模型分为三类:不可约表示、正则表示和标量化。基于信息不可约表示模型此类模型基于表示论中紧群的线性表示可以分解为一系列不可约表示的直积的理论。因此,在SE(3)群中构造了一个满足等变性的消息模型。例如,在TFN中:TFN层利用Clebsch-Gordan系数的特性构造一个参数,该参数对于任何旋转参数都与属于SO(3)的任何旋转操作等变。基于TFN结构扩展了很多工作,比如加入Attention机制,引入非线性Clebsch-Gordan系数等,但这些方法计算复杂度高,不可约表示只适用于特定群体。这限制了此类模型的表达能力。基于规范表示信息的模型另一类工作试图利用组规范表示来构造组卷积运算。这里的代表作李卷积(LieConv),通过Lifting操作将输入映射到组中的元素,然后使用PointConv完成组卷积的离散化计算。在我们的记法约定下,李卷积可以表示为:其中是群中的映射元素,log将群元素映射到对应的李代数上,就是一个MLP。通过这种构造,李卷积中h_i的更新实现了对任意李群及其离散子群的不变性。基于这种思想,LieTransformer引入了self-attention机制,进一步提升了模型的性能。基于李群正则表示的模型在组选择上更加灵活,但由于离散化和采样,需要在效率和性能之间做出权衡。同时,上述更新只考虑了标量信息h,很难直接扩展到几何信息x的更新,除非集成哈密顿网络等工作中的更新方法。除了基于群表示理论的标量化方法外,许多工作还采用了基于标量化的等变性建模方法。这类标量化方法首先将几何特征转化为一些不变的标量,然后利用MLP等网络结构得到一个标量变化,最后将这个变化加回到原来的几何特征上得到等方差。这种标量化方法最早由SchNet和DimNet提出,但只考虑了模型的不变部分。SphereNet建立在之前的工作之上,进一步考虑了跨消息传播网络的扭曲角变化。EGNN作为scalarization中的重要工作,提出了一个非常灵活的框架:其中,它是几何特征的scalarization,function是不同的MLP。通过关联几何信息和非几何信息消息,EGNN可以同时保证非几何和几何特征传播过程中的等变性。这种结构结合了物理知识,可以看作是对两个粒子的库仑力/重力的计算进行建模。GMN在EGNN的基础上扩展了模型可以描述的几何特征维度。在对坐标信息进行建模的同时,还可以引入更多的几何信息(如速度、加速度、角速度等),保证等变性。GemNet通过这种基于DimeNet的通用表示,将一些更丰富的几何特征,比如二面角,结合到消息传播的过程中。此外,存在一类基于不变标量和等变向量的乘积仍然是等变向量的观察构造等变消息传播的标量化方法。例如,PaiNN和EquivariantTransformer通过在不变的SchNet上用径向基函数建模原子的距离,将等变性质扩展到SchNet。以下是对等变图神经网络模型的总结:等变图神经网络应用由于能够更好地对几何信息进行建模,等变图神经网络被应用于从物理系统到化学物质的各类现实世界几何数据中,具有广泛的应用前景。应用范围。本综述简要介绍了其在物理系统、分子数据和点云数据中的应用。下表总结了现有等变神经网络的应用方向和数据集:复杂物理系统建模复杂物理系统的动力学建模一直是一个具有挑战性的课题。在物理系统中,存在带电粒子等物体,它们相互作用以通过基于物理定律的力产生轨迹。NRI的这项工作首次引入了n体模拟问题。n体系统包含多个由它们之间的库仑力驱动的带电粒子。n体问题的目标是在给定系统初始条件(坐标、速度和电荷)的情况下预测这些粒子的动力学轨迹。这个任务是E(3)等变的。SE(3)-Transformer和EGNN都展示了等变图神经网络在此任务上的潜力。在GMN中,进一步提出了一个更具挑战性的问题——带约束的n体问题,即当粒子之间存在约束(如链接或铰链)时,如何对粒子的轨迹进行有效预测。除了微观方向的数据,NRI和GMN还使用了人体动作捕捉的宏观数据来验证模型的有效性。分子建模同构神经网络的另一个重要应用方向是分子数据的建模。在分子数据中,原子的相互作用是由一系列复杂的化学和物理机制决定的。对于典型的分子数据,原子的非几何特征往往包括原子本身的一些特征,而几何特征则是原子的空间坐标和速度。原子之间的边缘由化学键构成或根据实际距离切断。经典的分子应用包括分子预测和分子生成。分子预测:具体来说,分子预测包括一些关于分子的性质和结构的预测任务。在分子预测领域,包括以下经典数据集。在小分子方面:QM9是一个经典的小分子数据集,具有12个量子特征预测任务。M17是在8个小分子上得到的动态轨迹数据集,其中还包含相应状态的能量和相互作用力等信息。ISO17是一个类似的分子动力学轨迹数据集,其中包含129种异构体的轨迹信息。OpenCatalyst2020(OC20)包含催化剂和底物催化过程的状态信息,其目标是在给定初始状态的情况下预测目标结构和相应状态的能量。在大分子方面:MDAnalysis是一套比较完整的蛋白质水平的分子动力学模拟数据。Atom3D是一个综合数据集,包含8个具有几何信息的分子预测任务,范围从小分子到RNA和蛋白质。分子生成:在分子生成领域,等变图神经网络常用于分子构象相关的生成。ConfGF和DSGM基于旋转平移不变GNN对评分函数进行参数化,并构建基于评分的构象生成模型。GeoDiff依赖于DenoisingDiffusionProbabilisticModel并基于具有等变保证的GNN构建模型。EquivariantFlow验证了基于等变核的NormalizingFlow的可行性。点云建模点云是对象的一种表示格式,它通过一组指定坐标的点来描述形状。在点云建模领域,包含了一些经典的数据和任务。ModelNet40和ScanObjectNN是两个经典的点云数据集,其任务是对物体进行分类。由于点云数据中不存在显式的点对点链接,等变图神经网络在对点云建模时,常常使用距离d作为阈值来构造点与点之间的边。TFN和SE(3)-Transformer在点云数据上都实现了与传统方法相比具有竞争力的性能。FutureProspects在系统总结了现有等变图神经网络的方法和应用之后。这篇综述还讨论了该领域一些潜在的未来发展方向:理论的完整性:与经典的图神经网络不同,等变神经网络仍然缺乏一系列关于表达性和泛化性的理论分析框架。一些现有的工作主要集中在讨论消息传播机制中存在的一般表示。然而,该模型的整体性质仍不清楚。如何构建一个完整的理论框架来指导模型的设计,是未来一个非常有趣的方向。大规模等变图神经网络,如上所述,基于群表示理论的方法存在计算复杂度高的缺点,限制了等变图神经网络在大规模数据上的应用。尤其是结合一些更复杂的结构,比如attentionmechanism,这个问题会更加严重。如何有效地简化现有模型并加速计算,使现有的等变图神经网络能够应用于大规模数据,是一个重要的方向。多层次结构建模:许多现实世界的系统都表现出复杂的层次结构。例如,有机分子由多个官能团组成,蛋白质由氨基酸组成。通过利用这些结构,我们可以设计模拟多粒度和多层次结构的系统。与现有的只有单层结构的消息传递范式相比。这种层次结构的等变模型可以更好地描述层次信息,提高模型的性能和泛化能力。新应用和数据:现有的等变图神经网络模型大多仅在规模和复杂性优先的系统上进行评估,例如模拟的N体系??统和小分子MD数据。未来,我们需要评估等态神经网络在一些更具挑战性的任务上的有效性,包括:更大数量的对象、更复杂的交互、更多样化的约束等。最近在蛋白质建模方面的工作是一些有趣的尝试。然而,由于数据收集困难和数据质量限制,尚未出现可以评估各种方法的综合数据集。未来如何将等变神经网络扩展到越来越复杂的领域来解决实际问题是一个非常有意义的方向。