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

等变图神经网络

SciencePedia玻尔百科
核心要点
  • 等变图神经网络(EGNNs)将旋转、平移等物理对称性直接嵌入其架构中。
  • 通过内在地理解几何学,EGNNs 实现了卓越的数据效率,无需进行大规模的数据增强。
  • 这些网络作用于标量、向量和张量等几何对象,并使用张量积等源于物理学的规则将它们组合起来。
  • EGNNs 在科学领域具有变革性的应用,能够在药物发现、材料科学、高能物理和气候建模中实现准确预测。

引言

标准的神经网络是强大的模式识别器,但它们缺乏基本的“物理常识”。对于一个典型的模型来说,一张旋转后的物体图片——或是一个模拟中旋转后的分子——是一个全新且无关联的数据。这迫使研究人员进行繁琐的数据增强过程,向网络展示无数个例子,以教会它一个基本的物理学原理:自然法则不因朝向而改变。这种低效率凸显了一个重大缺口:我们如何构建能够内在地理解并尊重物理世界基本对称性的人工智能?

本文探讨了解决方案:等变图神经网络(EGNNs),这类模型将几何与对称性的语言直接融入其计算结构中。通过这样做,它们超越了简单的模式匹配,走向一种植根于物理定律的推理形式。我们将首先深入探讨定义 EGNNs 的核心​​原理与机制​​,探索等变性与不变性等关键概念,以及赋予它们生命的数学工具,如张量积和球谐函数。随后,我们将通过对​​应用与跨学科联系​​的综述来探索该方法的深远影响,展示这些具备对称性意识的模型如何在药物发现、材料科学、高能物理和全球气候建模等领域开辟新前沿。

原理与机制

想象一下你在教一个孩子认识猫。你给他看一张猫坐着的照片,然后又给他看一张同一只猫侧躺的照片。“这两只都是猫,”你说。你必须为无数种姿势——倒立的、倾斜的、从后面看的——都这样做一遍。这个孩子因为缺乏三维物体的基本概念,不得不记住每一个视角。这就是标准神经网络的困境。对于一个计算机视觉模型来说,一张旋转后的图像只是一个完全不同的像素值矩阵。为了教会它旋转后的猫仍然是猫,我们必须在一个称为​​数据增强​​的过程中,费力地向它展示成千上万个例子。

现在,让我们从猫转向化学。一个水分子的能量并不取决于它在你的实验室里——或者在你的计算机内存中——是朝上、朝下还是朝向侧面。这是一个基本的物理学原理:支配自然的法则是相同的,无论你的位置或朝向如何。这就是​​物理定律的对称性​​。但标准的神经网络对此视而不见。对于你呈现给它的每一种可能的朝向,它都必须从头开始学习水分子的物理特性。这不仅仅是低效,更感觉是极其不智能的。它缺乏物理常识。我们如何才能将这种常识直接构建到我们人工智能的结构中呢?

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

要将物理学融入神经网络,我们首先需要一种精确的语言来描述对称性。让我们考虑三维空间中的刚体运动群组——所有可能的平移和旋转。这个群组被数学家和物理学家称为​​特殊欧几里得群​​,或 SE(3)SE(3)SE(3)。当我们对一个物体应用 SE(3)SE(3)SE(3) 变换时,其属性主要以两种方式响应。

想象一个指向北方的风向标。如果一阵西风吹来,风向标会旋转指向西方。它的朝向,一个向量,其变化方式与风向的变化完全耦合。这种性质称为​​等变性​​。如果一个函数在输入变换时,其输出也以相应、可预测的方式变换,那么该函数就是等变的。作用于分子中原子的力就是如此:如果你旋转分子,作用于每个原子上的力向量也会随之旋转。

现在,考虑一个数字仪表上显示的风速。它可能读作 151515 km/h。当风从北风转向西风时,读数仍然是 151515 km/h。这种性质称为​​不变性​​。如果一个函数在输入变换时其输出不发生改变,那么该函数就是不变的。我们水分子的总势能就是如此:它是一个单一的数值,无论分子如何在空间中翻滚,这个数值都保持不变。

不变性只是等变性的一种特殊情况,即输出的变换是“平凡的”——也就是说,它根本不改变。因此,最大的挑战在于设计出对于类向量属性本质上是 SE(3)SE(3)SE(3)-等变的,而对于类标量属性本质上是 SE(3)SE(3)SE(3)-不变的神经网络。

通往对称性的两条路径

历史上,研究人员主要通过两条路径来为化学和材料科学构建对称模型。

不变性路径

第一种方法在概念上很简单:如果你希望最终的预测是不变的,只需确保网络只看到不变的信息。以 SchNet 等模型为代表的这种架构,首先将原子几何结构转换为一组已经对旋转和平移免疫的特征。最明显的此类特征是任意两个原子之间的距离,∥rij∥\|\mathbf{r}_{ij}\|∥rij​∥。 然后,网络就变成一个标准的机器学习模型,学习这些距离与总能量之间的关系。

这是一个聪明的技巧。它保证了预测的能量是不变的。而且,作为微积分的一个优美推论,如果你将力定义为这个不变能量的负梯度(Fi=−∇riE\mathbf{F}_i = -\nabla_{\mathbf{r}_i} EFi​=−∇ri​​E),那么这些力就自动地、完美地具备了等变性!

但这条路径有一个显著的缺点。通过仅使用一维距离列表来描述三维结构,你丢弃了大量的几何信息。例如,仅通过观察分子的内部距离,你无法区分一个分子与其镜像(这种性质称为手性,在生物学中至关重要)。 你也很难描述那些本质上具有方向性或各向异性的现象,例如在催化剂复杂的阶梯状表面上的化学键合。 你正在迫使模型通过一个滤掉了所有朝向丰富性的针孔来看世界。

等变性路径:说几何的语言

这就引出了第二种更强大的方法:我们不回避方向性信息,而是拥抱它。这就是等变图神经网络(EGNNs)背后的哲学。其核心思想是允许网络内部的特征本身就是几何对象——不仅仅是普通的数字,而是标量、向量,甚至是更复杂的称为张量的对象。然后,网络学会在几何的语言中“思考”。

这场变换的交响乐是如何编排的?它依赖于群表示论中几个深奥而又被精妙地付诸实践的概念。

​​不可约表示 (irreps):​​ 这个听起来吓人的名字指向一个简单而强大的思想:根据物体在旋转下的行为对其进行分类。我们可以用一个整数 ℓ≥0\ell \ge 0ℓ≥0 来标记物体。

  • 类型为 ℓ=0\ell=0ℓ=0 的对象是​​标量​​。它在旋转时不会改变。
  • 类型为 ℓ=1\ell=1ℓ=1 的对象是​​向量​​。它有 3 个分量,以我们熟悉的方式进行旋转。
  • 类型为 ℓ=2\ell=2ℓ=2 的对象是​​二阶张量​​(如四极矩),它有 5 个分量,以一种更复杂但被完美定义的方式进行变换。 一个 EGNN 的特征是这些不可约表示的集合。

​​球谐函数:​​ 为了将几何信息输入网络,我们使用相对位置向量 rij\mathbf{r}_{ij}rij​ 来描述两个原子之间的关系。这个向量可以分解为两部分:它的长度(距离),是一个不变的标量(ℓ=0\ell=0ℓ=0);以及它的方向。这个方向可以由球面上的一组称为​​球谐函数​​ (YℓmY_{\ell}^{m}Yℓm​) 的函数优雅地描述。这些函数是朝向的天然“构建模块”,对于给定的 ℓ\ellℓ,每一组函数都作为一个 ℓ\ellℓ 类型的不可约表示进行变换。

​​张量积:​​ 这是计算的核心。网络如何组合信息——比如说,将一个原子上类型为 ℓ1\ell_1ℓ1​ 的特征与指向它的类型为 ℓ2\ell_2ℓ2​ 的几何方向结合起来?它使用一种称为​​张量积​​的运算。在被称为​​Clebsch–Gordan 分解​​的严格、源于物理的规则的支配下,这个运算精确地规定了可以形成哪些新类型的几何对象 (ℓout\ell_{out}ℓout​)。例如,组合两个向量(ℓ1=1,ℓ2=1\ell_1=1, \ell_2=1ℓ1​=1,ℓ2​=1)可以产生一个标量(ℓout=0\ell_{out}=0ℓout​=0,它们的点积)、另一个向量(ℓout=1\ell_{out}=1ℓout​=1,它们的叉积)和一个二阶张量(ℓout=2\ell_{out}=2ℓout​=2)。网络学习创造每种新对象的数量,但它受到几何学基本规则的约束。

​​宇称:​​ 除了旋转,我们可能还关心反射(镜像)。这由一个称为​​宇称​​的附加属性来处理。每个不可约表示都有一个宇称(偶或奇),组合它们的规则也包括宇称守恒。这使得 EGNNs 能够区分左手性和右手性分子,这是药物发现和分子生物学中的一个关键特征。

本质上,一个 EGNN 的消息传递层是一台被精妙约束的机器。它将几何对象作为输入,使用张量积这一不可改变的法则将它们与其邻域的几何形状相结合,并输出一组新的有效几何对象。每一层都尊重三维空间的内在对称性,不是因为它被训练成这样,而是因为它在物理上就不可能不这样做。

回报:更智能、更快速的科学

将物理定律直接构建到网络架构中会带来深远的影响。

首先,网络变得极其​​数据高效​​。它不需要看到一个分子的一千种不同朝向才能理解其性质;看一次就足够了。内置的等变性使其能够即时泛化到任何其他朝向。这意味着用分子的旋转副本进行数据增强变得完全多余——它无法为已经掌握旋转语言的模型提供任何新信息。

其次,这种架构的优雅性使得强大的​​多任务学习​​成为可能。我们可以设计一个单一网络,共享一个丰富的、等变的内部表示,以同时预测多种属性。例如,一个模型可以预测药物的标量结合亲和力(一个不变属性),同时预测作用于其原子的向量力(一个等变属性)。这是通过设置独立的“读出”头来实现的,这些头分别利用适当的特征:不变的标量(ℓ=0\ell=0ℓ=0)特征用于亲和力,等变的向量(ℓ=1\ell=1ℓ=1)特征用于力。这种方法不仅高效,而且在物理上也是一致的,特别是当力是作为学习到的势能的梯度导出时。

通过拥抱对称性与物理学之间的深刻联系,等变神经网络代表了一种范式转变。它们超越了模式识别,走向一种植根于宇宙基本结构的计算推理形式。它们不仅仅是从数据中学习,更是在自然法则的指引下学习。

应用与跨学科联系

在了解了等变图神经网络的原理和机制之后,我们可能会倾向于将它们视为一种优美但抽象的数学构造。事实远非如此。在科学和工程领域,我们不断面临一个基本现实:自然法则对我们的观察视角漠不关心。实验的结果不取决于我们实验室的朝向或坐标系的原点选择。这一被称为对称性的原则,不仅仅是一种哲学偏好,更是对任何有效物理理论的深刻而强大的约束。等变神经网络正是这一原则在机器学习语言中的体现。它们不仅仅是一个巧妙的技巧,更是一种更忠实地模拟物理世界的方式。

通过将对称性直接构建到其架构中,这些网络在众多学科中成为了极其有效的工具。它们学习得更快,泛化能力更强,并且其产生的结果在设计上就保证了物理一致性。现在,让我们来探索这一领域,从分子的微观世界到我们整个地球的宏大尺度,看一看等变性这一统一的思想是如何开启新的发现前沿的。

分子与材料的世界

等变网络最自然的应用领域或许是分子科学,在这里,原子几何决定了一切。其核心对称性是三维空间的对称性,即被称为 SE(3)\mathrm{SE}(3)SE(3) 的旋转和平移群。一个分子的物理性质——其能量、稳定性、反应性——取决于其原子的相对排列,而不是其在空间中的绝对位置或朝向。

一个典型的例子来自至关重要的药物发现领域。一项关键任务是预测潜在药物分子(配体)与靶点蛋白之间的结合亲和力。这种决定药物效力的亲和力是一个单一的标量值——一种能量。一个等变 GNN 可以接收蛋白质-配体复合物的三维结构并预测该能量。由于网络被构建为不变的,无论复合物如何旋转或平移,其预测结果都将完全相同。这不仅仅是学术上的讲究,更是一个实用模型的关键要求。旧方法仅依赖原子间距离,虽然能捕捉到部分不变性,但却丢弃了关于化学键和相互作用的关键方向性信息。相比之下,等变网络处理和更新明确表示这种方向性几何的向量和张量特征,同时保证最终的单一能量值正确地尊重了底层的对称性。

这种方法的力量超越了单一的能量值。在许多情况下,结合过程是由在界面处形成精细氢键网络的“桥接”水分子介导的。忽略它们会导致预测不准确。等变 GNNs 可以被训练来预测这些关键水分子的最佳位置和能量贡献,从而有效地将一个初始的、粗糙的对接姿态优化为物理上真实可信的结构。网络通过处理来自邻近原子的信息来学习放置水分子。每条信息都是一个向量,从邻近原子指向水分子,并由一个学习到的关于距离和化学特征的函数进行缩放。将这些等变向量相加提供了一个方向性的更新,将水分子推向其最有利的位置——如果整个系统被旋转,这个预测会正确地进行变换。

这种预测结构化、几何量的能力是等变模型最深刻的方面之一。它们不限于标量输出。例如,在材料科学中,我们希望预测原子如何响应电场。这不仅需要知道它们的电荷(标量),还需要知道它们的偶极矩(向量)、四极矩(二阶张量)以及极化率(也是二阶张量)。一个先进的等变 GNN,使用群论的形式化语言以及球谐函数和 Clebsch-Gordan 系数等工具构建,可以学习同时预测所有这些量。每个输出都由具有数学上正确的张量特性的网络特征构建,确保当分子旋转时,预测的偶极矩向量也随之旋转,而预测的极化率张量则完全按照一个二阶张量应有的方式进行变换。这甚至可以扩展到学习量子力学的基本算符。一个多头 GNN 可以被训练来输出分子的整个哈密顿矩阵,由此可以计算出势能面和非绝热耦合,这些正是驱动视觉和光合作用等光化学反应的过程。

从性质到动力学,下一个合乎逻辑的步骤是模拟材料随时间的行为。这是分子动力学(MD)的领域,我们计算每个原子上的力,并用它们来推动系统在时间上演化。人们的梦想是拥有一个“通用力场”,它既有量子力学般的精度,又像简单的经典模型一样快。等变 GNNs 正在使这一梦想成为现实。通过学习势能面——一个关于所有原子位置的标量函数——它们只需通过对预测能量取解析梯度就能提供力。这保证了力是保守的(能量守恒),并且至关重要的是,力是连续的,这对于稳定和准确的模拟至关重要。这些机器学习原子间势(MLIPs)可以捕捉到在金属表面催化反应中控制键断裂和形成的复杂多体效应,这一过程远非传统的、简化的力场所能及。此外,这些模型可以在“主动学习”循环中使用。一个由不同等变势组成的委员会可以用来估计模型自身的不确定性。在不确定性高的地方,我们可以指挥超级计算机执行一些昂贵的量子计算以生成新的训练数据,从而在最需要的地方迭代地优化势函数。

从原子到工程结构

等变性还在原子的微观世界与工程师用于设计汽车、建筑和飞机的连续介质力学宏观世界之间架起了一座强大的桥梁。在这个宏观世界中,一个核心量是柯西应力张量 σ\boldsymbol{\sigma}σ,它描述了材料内部的内力。我们如何能从底层晃动原子的混沌状态中确定某一点的应力?

一个等变框架提供了一个有原则的答案。我们可以训练一个模型,将原子的局部排列映射到相应的应力张量。这个映射必须遵守两个基本的力学定律。首先,应力张量必须是对称的。其次,这种关系必须是客观的或坐标系无关的,这意味着如果我们刚性地旋转这片原子,得到的应力张量必须是相同的物理量,只是在新的、旋转后的坐标系中表示。这正是张量等变性的精确定义。一个等变 GNN 可以被设计为接收一团原子作为输入,并输出一个二阶张量,该张量通过其构造保证了对称性和客观性。这使得我们能够创建数据驱动的材料模型,这些模型直接从原子模拟中学习复杂的非线性行为,同时严格遵守连续介质力学的基础原则。

从粒子径迹到全球气候

等变性原理是如此基础,以至于其应用延伸到我们研究的最小和最大尺度。在高能物理学中,大型强子对撞机(LHC)的实验会产生大量的粒子,它们在巨大的探测器中留下离散的“击中点”。挑战在于从这个庞大的点云中重建带电粒子的弯曲轨迹或“径迹”。这是一个宇宙级的连点成线游戏,但规则必须是物理的。径迹的存在是一个物理事实,与探测器在其大厅中的放置方式无关。

一个等变 GNN 可以学习解决这个问题。网络可以观察成对的击中点,并根据它们的相对位置,学习一个不变的分数,该分数代表它们属于同一径迹的概率。通过组合三个击中点的分数,并为高曲率添加一个基于几何的惩罚项(直线是“默认”径迹),该算法可以有效地识别有希望的径迹种子。GNN 内部的位置更新是等变的,能正确处理三维几何,而最终的评分则使用不变的输出来做出独立于坐标系的决策。

在另一个极端,是气候建模的巨大挑战。大气环流模型将大气和海洋离散化到网格上。然而,许多关键的物理过程,如云的形成,发生在比单个网格单元更小的尺度上。这些“次网格”过程必须被近似,或“参数化”。一个物理的参数化方案无论在地球上的哪个位置,其行为都应相同。领域的对称性决定了学习到的参数化方案所必需的架构。

对于一个简单的、平坦的矩形网格上的模型,其潜在的对称性是平移对称性。网格单元 A 中的物理过程应与网格单元 B 中的物理过程相同。标准的卷积神经网络(CNN)通过使用共享权重的卷积核,内建了这种平移等变性,使其成为一个自然的选择。但地球不是平的。如果我们将球面投影到一个矩形的经纬度网格上,一个标准的 CNN 将会学习到扭曲的、不符合物理的假象,将两极与赤道区别对待。为了尊重球体真正的旋转对称性 SO(3)\mathrm{SO}(3)SO(3),我们需要一个更复杂的架构,比如球形 CNN 或在更均匀的二十面体网格上运行的 GNN。架构的选择不是为了方便,而是关于尊重我们试图解决的问题的基本几何学的一个深刻声明。

从单个分子到整个地球,我们得到的教训是相同的。通过拥抱自然的对称性,等变神经网络为科学发现提供了一个更有原则、更稳健、更强大的框架。它们让我们能够构建使用物理学母语的模型,在这种语言中,法则是普适的,而我们人为选择的坐标系则理所当然地无关紧要。