try ai
科普
编辑
分享
反馈
  • 等变图神经网络 (GNNs)

等变图神经网络 (GNNs)

SciencePedia玻尔百科
核心要点
  • 等变 GNN 直接将旋转、平移等基本物理对称性整合到其架构中,确保输出随输入进行可预测的变换。
  • 模型通过两种主要策略实现 E(3) 等变性:一是仅使用距离、角度等不变特征;二是通过保持对称性的操作处理向量、张量等几何特征。
  • 通过嵌入对称性,这些模型的数据效率大大提高,并能保证其预测与能量守恒等物理定律保持一致。
  • 等变 GNN 通过提供更准确、更可靠的模拟,正在分子动力学、材料设计、粒子物理学和机器人学等领域引发科学发现的变革。

引言

物理定律拥有一个深刻而优雅的特性:它们是对称的,无论观察者的位置或朝向如何,都保持不变。当我们构建物理世界的计算模型时,我们面临一个关键选择——要么希望模型能从海量数据中学习到这些对称性,要么将它们直接嵌入到模型的设计中。等变图神经网络(GNNs)采用了后者,这种更强大的方法,创造出能用与自然本身相同的几何语言“思考”的模型。本文旨在解决构建尊重基本物理原理的人工智能这一挑战,从而实现更高的效率和准确性。在接下来的章节中,我们将首先探讨等变性的核心“原理与机制”,区分其与不变性,并概述构建对称性感知的 GNN 的关键策略。然后,在“应用与跨学科联系”部分,我们将见证这些有原则的模型如何正在彻底改变从化学、材料科学到高能物理和机器人学的各个领域。

原理与机制

假如你正试图描述自然法则,你会很快发现一个深刻而美丽的真理:这些法则不依赖于你的视角。无论你是在伦敦还是东京,在今天还是明天,面朝北方还是南方进行实验,其底层的物理学都保持不变。这个宏伟的思想,即对称性原理,不仅是一种哲学上的雅趣,更是我们构建对宇宙理解的深刻而有力的指南。

当我们为物理系统(无论是一个分子、一个星系,还是天气)构建计算模型时,我们有一个选择。我们可以要么期望我们的模型从海量数据中学习到这些基本对称性,要么将对称性直接构建到模型本身的结构中。后一种方法,即​​等变性​​(equivariance)之路,不仅更优雅,而且功能更强大、效率也高得多。它是一种物理原理在架构上的体现。

对称性的语言:不变性与等变性

让我们首先明确我们的术语。当我们谈论三维空间的对称性时,我们指的是​​欧几里得群 E(3) ​​,它包含了所有可能的刚体运动:平移(移动而不转动)、旋转(转动而不移动)和反射(照镜子)。一个尊重这些对称性的模型必须表现出以下两种行为之一:

  • ​​不变性​​(Invariance):当输入被变换时,模型的输出完全不发生改变。想象一个水分子的总势能。它是一个单一的数字,一个标量。这个能量取决于氢原子和氧原子的相对排列,但与分子是在你的实验室里还是在月球上无关,也与它指向哪个方向无关。能量在平移和旋转下是​​不变的​​。

  • ​​等变性​​(Equivariance):当输入发生变换时,模型的输出会以一种可预测的方式相应地变换。思考作用于同一个水分子中各个原子的力。每个力都是一个向量,既有大小也有方向。如果你旋转分子,你会期望力向量也随之旋转。从分子的几何结构到其受力的映射是​​等变的​​。输出与输入和谐共舞。

因此,我们的挑战就是构建一个能够内在地理解这种舞蹈的学习机器——一个图神经网络(GNN)。

从置换到几何:一个热身

在处理完整的三维空间几何之前,让我们先考虑一个更简单的对称性。想象一下你正在分析欧洲核子研究中心(CERN)探测器记录的一次粒子碰撞事件。这个事件只是一堆或一组粒子。你列出它们的顺序是完全任意的。一个物理结论,比如碰撞的总能量,不应依赖于这种任意的标签。这就是​​置换不变性​​(permutation invariance)原理。

我们如何构建一个自动具备置换不变性的神经网络呢?一个优美而简单的方案是像 Deep Sets 这类模型的基础:

  1. 取集合中的每一项(每个粒子的特征向量)。
  2. 将每一项输入到一个相同的神经网络中,我们称之为 ϕ\phiϕ。
  3. 使用一个不关心顺序的操作(如求和或平均)来聚合结果。
  4. 将这个单一的聚合结果输入到最后的神经网络 ρ\rhoρ 中,得到你的答案。

这个结构,f(X)=ρ(∑iϕ(xi))f(X) = \rho(\sum_{i} \phi(x_i))f(X)=ρ(∑i​ϕ(xi​)),通过其设计本身就保证了置换不变性。打乱输入只会打乱求和项的顺序,但这并不会改变结果。

这是一个深刻的架构思想。我们不仅仅是训练了一个在所见过的数据上恰好表现出不变性的模型;我们构建了一个对于任何可能的输入都不可能不是不变的模型。这就是等变工程的精神。图神经网络自然地扩展了这一思想,以产生逐粒子的输出,其中对称聚合发生在消息传递步骤中,确保粒子 i 的输出对应于粒子 i 的输入,这一属性被称为​​置换等变性​​(permutation equivariance)。

几何和谐的两大策略

现在,让我们回到完整的三维几何。我们的 GNN 节点不再是抽象的粒子,而是空间中具有位置 ri\mathbf{r}_iri​ 的原子。我们如何尊重 E(3) 群呢?第一个对称性,平移,处理起来出奇地简单。在空旷空间中的物理相互作用取决于相对位置,而非绝对位置。因此,我们设计的 GNN 只会看到​​相对位置向量​​ rij=ri−rj\mathbf{r}_{ij} = \mathbf{r}_i - \mathbf{r}_jrij​=ri​−rj​。由于全局平移会给 ri\mathbf{r}_iri​ 和 rj\mathbf{r}_jrj​ 加上相同的向量,它们的差值保持不变,我们的模型因此自动变得平移不变。

真正的挑战在于旋转。于此,两种优雅的策略应运而生。

不变性之路:一个没有方向的世界

一种方法是从一开始就构建一个对方向“视而不见”的模型。我们可以构建本身就具有旋转不变性的特征,并将它们输入到一个标准的 GNN 中。这些特征是什么呢?

  • 两个原子间的距离,∥rij∥\|\mathbf{r}_{ij}\|∥rij​∥。
  • 两个化学键之间的夹角,θijk\theta_{ijk}θijk​,可以通过点积 cos⁡θijk=r^ji⋅r^jk\cos \theta_{ijk} = \hat{\mathbf{r}}_{ji} \cdot \hat{\mathbf{r}}_{jk}cosθijk​=r^ji​⋅r^jk​ 求得。
  • 两个原子平面之间的二面角,可以通过相对位置向量的点积和叉积求得。

所有这些——距离、角度、二面角——都是​​标量不变量​​。它们是在旋转下值保持不变的数字。一个只看到这些不变特征的模型,就像许多经典势场和一些机器学习模型(如基于 SOAP 的高斯近似势)一样,会自然地产生一个不变的输出,非常适合用来预测总能量 EEE。

但是力 Fi\mathbf{F}_iFi​ 怎么办呢?它们必须是等变向量!这是否意味着不变性之路走到了死胡同?这其中蕴含着一个数学魔术般的时刻。矢量微积分的一个基本定理指出:​​一个不变标量场的梯度是一个等变向量场​​。这意味着,如果我们有一个模型能够正确地将总能量 EEE 预测为原子位置的可微函数,我们就可以通过取负梯度来计算力,即 Fi=−∇riE\mathbf{F}_i = -\nabla_{\mathbf{r}_i} EFi​=−∇ri​​E。根据这一定理,得到的力向量必然是完全 E(3) 等变的。那些看似被丢弃的方向信息,通过微分这一行为被奇迹般地恢复了。

等变性之路:教向量学会舞蹈

另一种策略更为直接。我们不再让模型对方向“视而不见”,而是教它几何规则。我们构建一个网络,其中的特征本身不仅仅是数字,而是懂得如何旋转的几何对象。

在 E(3) 等变 GNN 中,与一个原子相关联的特征可能是一组标量(0 型,不变)、向量(1 型,会旋转),甚至更高阶的张量(2 型等,具有更复杂的旋转规则)。GNN 的消息传递层就是根据物理和群论的法则来组合这些几何对象的操作。

例如,要用来自邻居 j 的信息来更新原子 i 上的特征,我们不能简单地将它们的笛卡尔向量分量拼接起来,然后输入到一个标准的多层感知机(MLP)中。MLP 将其输入视为一个简单的数字列表;对向量的 xxx、yyy 和 zzz 分量独立应用像 ReLU 这样的非线性函数,会破坏其几何特性,从而打破等变性。

相反,我们必须使用尊重几何的操作。我们可以通过两个向量的​​点积​​来形成新的标量(不变量)。我们可以通过​​叉积​​来形成新的向量。构建等变更新的一个强大而通用的方法是,将新向量构造为现有等变基向量的线性组合,其中系数本身是由 MLP 计算出的不变标量。例如,从 jjj 到 iii 的一条消息可能是一个如下的向量: mij=α1(invariants) ui+α2(invariants) uj+α3(invariants) r^ij\mathbf{m}_{ij} = \alpha_1(\text{invariants}) \, \mathbf{u}_i + \alpha_2(\text{invariants}) \, \mathbf{u}_j + \alpha_3(\text{invariants}) \, \hat{\mathbf{r}}_{ij}mij​=α1​(invariants)ui​+α2​(invariants)uj​+α3​(invariants)r^ij​ 在这里,ui\mathbf{u}_iui​ 和 uj\mathbf{u}_juj​ 是已有的向量特征,r^ij\hat{\mathbf{r}}_{ij}r^ij​ 是它们之间的方向向量,而标量系数 αk\alpha_kαk​ 是关于距离和点积等不变量的学习函数。这种构造保证了 mij\mathbf{m}_{ij}mij​ 会正确地旋转。

这些模型中最复杂的版本,应用于理论化学和物理学中,使用量子力学的语言将其形式化。它们将特征表示为旋转群的​​不可约表示​​(或称“irreps”),由角动量数 lll 索引。然后,它们使用​​张量积​​和 ​​Clebsch-Gordan 系数​​来组合这些特征——这正是用于将在原子中电子的角动量相加的同一套数学工具。最终的能量(一个 l=0l=0l=0 的标量)和力(l=1l=1l=1 向量的集合)随后通过将最终的丰富几何特征投影到所需的输出类型上来读出。

回报:为何这套精美的机制至关重要

将这种对称性构建到网络的架构中不仅仅是一种审美上的追求;它具有巨大的实际意义。

  • ​​数据效率​​:一个等变模型不需要学习什么是旋转。它已经知道了。当它从单一的分子构型中学习时,它自动理解了该构型的所有无限多个旋转和平移副本的物理性质。与非对称模型相比,这极大地减少了达到高精度所需的训练数据量。

  • ​​物理一致性​​:通过构造,模型的预测保证遵守物理学的基本对称性。你永远不会得到旋转一个分子会改变其预测能量或导致力不能正确旋转这样的非物理结果。

  • ​​更智能的科学发现​​:在​​主动学习​​等前沿应用中,算法必须智能地决定要运行哪些新的模拟,此时等变性是一种超能力。一个等变模型的不确定性估计也是不变的。它认识到先前见过的结构的旋转版本并非“新”的,并且不会浪费昂贵的计算资源来重新计算它已经知道的东西。它将搜索集中在广阔化学空间中真正新颖且信息丰富的区域,从而加速发现。

从本质上讲,通过将深刻的对称性原理直接嵌入到我们神经网络的结构中,我们创造出的模型不仅更准确、更高效,而且还能用与自然本身相同的几何语言“思考”。

应用与跨学科联系

在探索了等变性的基本原理之后,我们现在来到了激动人心的部分:见证这些思想在实践中的应用。欣赏一个概念的数学优雅是一回事,而亲眼目睹它在广阔的科学和工程领域解决实际问题的力量则完全是另一回事。正如我们所见,对称性原理不仅仅是一种审美偏好;它是一个强大的约束,一盏指路明灯,帮助我们构建不仅更准确,而且数据效率更高、物理意义更强的模型。

现在,让我们开启一次跨学科之旅,看看配备了对称性语言的等变图神经网络(GNNs)如何彻底改变我们理解和与世界互动的方式,从最小的分子到最大的宇宙事件。

化学家的梦想:从零开始构建分子

也许 E(3) 等变模型最天然的归宿是分子的世界。毕竟,分子是典型的三维结构——存在于空间中的原子集合,支配它们相互作用的物理定律不依赖于你的视角。如果你模拟一个水分子,无论它是在你帕萨迪纳的实验室里,还是在环绕木星的宇宙飞船中,结果都应该是一样的,而且当然不应取决于你观察它的方式。

化学和材料科学中的一大挑战是计算原子系统的势能面。如果你知道了任何给定原子排列下的能量,你几乎可以推导出其他所有信息。例如,作用在每个原子上的力就是能量的负梯度(最陡下降方向)。这些力反过来又允许你进行分子动力学模拟——实时观察蛋白质折叠、药物与靶点结合,或化学反应的展开。

几十年来,这些计算都是耗费巨大的量子力学方法的领域。但是,如果我们能直接从数据中学习这种能量-结构关系呢?这正是等变 GNN 大放异彩的领域。网络以原子结构(一个带有三维位置的原子图)为输入,并被训练来预测一个关键的数字:系统的总势能。

这种方法的美妙之处有两点。首先,因为能量是一个标量,它必须在旋转和平移下是不变的。一个等变 GNN 架构可以被设计成产生一个保证不变的标量最终输出。其次,也是真正神奇的部分,一旦网络学会了标量能量函数,我们就可以免费获得每个原子上的矢量力,只需对网络输出关于输入原子位置取解析梯度。这是通过训练网络时使用的同一个反向传播机制自动完成的!这个优雅的技巧确保了学习到的力在构造上是能量守恒的——这是一个幼稚的模型可能会违反的基本物理定律。

但应用并不仅限于力和能量。许多关键的分子性质不是简单的标量或向量,而是由更复杂的数学对象——张量——来描述。一个完美的例子是分子的​​极化率​​(polarizability),它描述了分子的电子云在外部电场作用下的形变方式。这个性质由一个二阶张量 α\boldsymbol{\alpha}α 表示,这是一个 3×33 \times 33×3 矩阵,告诉你一个方向的电场如何在另一个方向上引起偶极矩。

如果你旋转分子,这个极化率张量必须随之以一种特定的方式旋转(α↦RαR⊤\boldsymbol{\alpha} \mapsto \mathbf{R}\boldsymbol{\alpha}\mathbf{R}^\topα↦RαR⊤)。一个只使用旋转不变特征(如原子间距)的 GNN 永远无法预测极化率张量的方向。它已经丢掉了解决问题所必需的方向信息。为了预测一个张量,GNN 本身必须用张量的语言来“思考”。E(3) 等变 GNN 正是这样做的。它们的构建使用了本身就是向量和张量的特征,这些特征以一种严格尊重其变换性质的方式被传递和处理。这确保了网络不仅能正确预测一个分子会极化,还能预测其响应在三维空间中如何取向。

材料科学家的水晶球:从分子到材料

从单个分子扩展到更大的尺度,我们进入了材料科学的领域。无数原子的集体相互作用如何产生材料的宏观性质,如其刚度或强度?在这里,对称性同样至关重要。​​材料框架无关性​​(material frame indifference)或客观性(objectivity)原理是连续介质力学的基石。它指出,材料的本构定律——即关联形变与应力的规则——必须独立于观察者的参考系。这恰恰是 E(3) 等变性原理的另一种表现形式。

等变 GNN 提供了一个强大的工具,以弥合原子世界和连续介质之间的鸿沟。想象一下,直接学习从一个小的原子邻域到该点宏观柯西应力张量的映射。E(3) 等变 GNN 是完成这项粗粒化任务的完美工具。它可以将原子的三维构型作为输入,并输出一个应力张量,该张量会随着底层原子结构的旋转而正确旋转,从而通过构造满足客观性原理。

当我们考虑晶体材料时,情况变得更加有趣。与气体或液体(各向同性,即在所有方向上看起来都一样)不同,晶体具有特定的内部结构——晶格——这赋予了它优先方向。例如,金刚石晶体的强度取决于你施加压力的方向。这些材料不是完全旋转对称的;它们只在一组特定的、有限的旋转和反射下对称,这些对称操作构成了它们的​​点群​​。例如,食盐晶体具有立方对称性,而苯分子具有 D6hD_{6h}D6h​ 群的六边形对称性。

值得注意的是,等变 GNN 的框架可以扩展到处理这些特定的材料对称性。通过采用更高级的群表示论技术,可以构建一个不是对所有可能的三维旋转都等变,而是专门对所建模材料的晶体学点群等变的 GNN。这使我们能够学习到高度精确、数据驱动的各向异性材料模型,从原子数据中捕捉它们复杂的方向性行为。

物理学家的眼睛:看见无形

几何等变性的用途远远超出了原子和材料的有形世界。考虑一下实验高能物理学的巨大挑战。在像大型强子对撞机(LHC)这样的粒子对撞机中,质子以接近光速的速度相互碰撞,产生大量向四面八方飞出的新粒子。一个巨大而复杂的探测器将这些粒子的通道记录为一系列“击中点”——三维空间中的点。物理学家的工作就是扮演侦探,将这些点连接起来,重建原始粒子的轨迹,即“径迹”。

这本质上是一个三维点云中的模式识别问题。但是,支配粒子轨迹的底层物理定律是旋转和平移对称的。因此,一个好的径迹寻找算法也应该是这样的。一个等变 GNN 可以在模拟数据上进行训练,以识别类似径迹的击中点模式。它学会给那些位于平缓曲线上的三重击中点(这是带电粒子在磁场中弯曲的特征)赋予高分。因为 GNN 是等变的,所以无论碰撞的方向如何,它都能找到这些径迹,使其成为解读这些极端实验结果的强大而高效的工具。

而且我们可以将对称性原理推向更远。LHC 的物理学不仅受三维空间对称性的支配,还受爱因斯坦狭义相对论更深层次的对称性支配,这由洛伦兹群描述。这个群不仅包括旋转,还包括“boosts”(或称“推助”)——即不同匀速运动的观察者之间的变换。我们能否构建一个尊重这个更大对称群的网络?答案是肯定的。通过构建专门操作洛伦兹不变量(从粒子的四维矢量导出)的网络,物理学家们正在设计通过构造即为洛伦兹等变的喷注标记算法。这确保了他们的预测与相对论的基本原理相一致,这是物理学与机器学习真正美丽的结合。

工程师的工具箱:从模拟到现实

等变性原理并不仅限于基础科学的深奥领域;它们在工程学中具有深远的实际意义。例如,在​​机器人学​​中,一个常见的任务是判断机器人对物体的抓取是否稳定。这种稳定性是抓取几何的内在属性,不应取决于机器人在哪里或物体在工作空间中如何定向。

我们可以将一次抓取建模为一个图,其中节点是物体表面的接触点。然后可以训练一个等变 GNN 来预测稳定性得分。一个非等变模型,当面对同一物体以不同方向呈现时,可能会给出完全不同且错误的预测。相比之下,一个专门构建的 SE(3) 不变模型每次都会给出正确的预测,因为它的架构本身就尊重问题的底层几何。这展示了编码已知对称性的巨大实用价值:它能产生更可靠、需要更少数据且泛化能力更强的模型。

最后,GNN 与传统的计算工程工具之间存在着深刻而强大的联系。几十年来,工程师和科学家通过在网格上求解偏微分方程(PDEs)来模拟物理现象——从机翼上的气流到地震波在地球中的传播。像有限元法(FEM)这样的方法通过定义一个局部“模板”(stencil)来工作,该模板将一个点上场的值与其直接邻居的值联系起来。

事实证明,一类特定的 GNN 可以被看作是这些数值方法的直接、可学习的模拟物。这样一层 GNN 执行的操作在数学上等同于一个学习到的 FEM 模板。它遵循相同的核心原则:局部性(消息只在相邻节点间传递),且其系数取决于网格的几何形状。通过保持常数场不变,它们正确地模仿了像拉普拉斯算子这样的算子的行为。这一见解是变革性的。它意味着我们不仅可以将 GNN 用作黑箱预测器,还可以将其用作物理启发的模拟器,直接从数据中学习求解 PDE,从而可能将模拟速度加快几个数量级。

从分子到材料,从夸克到机器人抓取,信息是明确的。自然界的对称性不是一个需要用数据增强来平均掉的麻烦,而是一个需要拥抱的深刻原理。通过将这些对称性构建到我们机器学习模型的结构中,我们创造出的工具不仅更强大,而且能用与它们试图描述的宇宙相同的语言进行交流。