无约束机器学习模型如何学习物理对称性

《PRX IntelligenceNEW》:How Unconstrained Machine-Learning Models Learn Physical Symmetries

【字体: 大 中 小 】 时间:2026年10月28日 来源:PRX IntelligenceNEW

编辑推荐:

  生成完全满足相应物理量基本对称性的预测这一需求,深刻影响了用于物理模拟的机器学习(ML)模型的发展。在许多情况下,模型是通过受限的数学形式构建的,以确保对称性被精确地强制执行。然而,人们常常发现,不遵守旋转对称性的无约束模型具有竞争性的性能,并且能够通过简单的数据增强策略学习到高

  生成完全满足相应物理量基本对称性的预测这一需求,深刻影响了用于物理模拟的机器学习(ML)模型的发展。在许多情况下,模型是通过受限的数学形式构建的,以确保对称性被精确地强制执行。然而,人们常常发现,不遵守旋转对称性的无约束模型具有竞争性的性能,并且能够通过简单的数据增强策略学习到高精度的近似等变行为。在本文中,我们引入了严格的度量标准来衡量此类模型中学习表征的对称性内容,并评估输出满足等变条件的准确性。我们将这些度量应用于两个在装饰点云上运行的无约束Transformer模型(一个用于原子模拟的图神经网络,以及一个用于粒子物理的PointNet风格架构),以研究对称性信息如何跨架构层处理以及如何在训练过程中被学习。基于这些见解,我们建立了一个严格的框架来诊断机器学习模型中的谱失效模式。通过这种分析,我们证明可以通过策略性地注入最小必需的归纳偏置来实现卓越的稳定性和准确性,在保持无约束架构的高表达性和可扩展性的同时,保证物理保真度。
I. 引言

对称性是现代物理学的基石。它们与守恒定律的深刻联系体现在诺特定理中,这一原则长期以来指导着将经验现象形式化为理论框架,并作为新理论发展的基础 [1,2]。物理科学的不同领域以不同的对称群为特征,例如分子力学中的正交群 [3],高能物理中的洛伦兹群 [4,5],以及量子力学和量子色动力学中的特殊酉群和 [6–8]。因此,在开发物理数据驱动模型时,结合物理对称性常被视为最自然的选择。这种方法推动了物理科学中机器学习(ML)方法的快速和富有成效的发展,其中原子模拟作为一个非常成功的应用领域脱颖而出 [3,9–13]。然而,确保模型严格保持群等变性,即保证输出在输入上的群作用下可预测地变换,会施加僵化的架构约束,这些约束可能计算代价高昂,并限制模型的表达性 [14,15]。

相反,主流计算机科学和机器学习方法已发展为最大化架构表达性和效率,让特定领域的归纳偏置直接从数据中学习。最近,人们对将这些“无约束”模型应用于物理和化学的兴趣日益浓厚,显著例子包括从生物分子结构预测(其中 AlphaFold 3 放宽了先前版本的一些严格等变约束 [16])到使用类似 PointNet 的架构 [20] 对等离子体晶体 [17]、星系形态 [18] 和粒子轨迹 [19] 进行分类。通过放宽严格的等变约束,此类模型可以获得拟合能力和计算效率,但必须从数据中推断相关的变换规律,通常通过对对称群进行增强 [21]。最近的工作 [22,23] 表明,与基线模型准确性相比,输出级等变误差通常很小,相关分析表明,在几个分子任务中,使用旋转增强训练的对称不可知模型可以在训练期间迅速减少这些误差 [24]。

在原子模拟领域,无约束模型作为显式等变模型成功的替代方案正日益受到关注,为量子力学提供了快速、准确且可迁移的代理模型。机器学习原子间势(MLIPs),即从原子位置和化学类型预测势能、原子受力和晶胞应力,是这方面最活跃的领域之一。然而,对于这些模型,近似对称性不能仅通过平均测试准确性来评估。原子结构的旋转或反转应引发矢量和张量预测(如受力和应力)的相应变换,残余的不一致性可能影响基于这些预测的模拟。最近的工作 [13,23,25,26] 表明,无约束模型在基准测试中可以匹配或超越不变和等变架构,同时通过可在当前加速架构中高效实现的操作提供卓越的可扩展性。无约束机器学习原子间势的能力超越了基准测试,在复杂材料科学模拟 [23] 和高通量筛选 [27,28] 等现实场景中表现出色。

更近期,无约束架构已被应用于超越不变预测(例如能量)的任务,扩展到矢量(例如力 [29,30])和更高秩张量目标 [31–33]。这些任务要求显著更高,要求模型通过数据增强学习复杂的变换规律。无约束架构似乎能够应对这些挑战,但可能需要更长的训练过程来采样等变模型通过设计编码的群轨道 [28]。

鉴于有经验证据表明这些模型可以学习对称性,自然会引发以下问题:这些对称性是在何时以及如何被学习的,是在整个模型架构中还是在训练过程中?这种知识能否推动设计性能更好的架构?我们引入度量标准来量化模型输出违反等变条件的程度,并将模型中编码的信息按照相关群的不可约表示进行分解,并利用它们定量分析对称性信息在架构中的流动及其在训练过程中的演变。然后我们将证明,实际上,深入模型“黑箱”内部并理解架构每个层级上对称性和数据的相互作用,可以让人对哪些归纳偏置是必需的、哪些是多余的做出明智决策,最终提高模型性能。

II. 量化机器学习模型中的对称性

我们关注一类在装饰点云上运行以预测物理相关量的机器学习模型 [见图1(a)],尽管类似的考虑适用于更一般的情境,例如三维(3D)形状识别或生成建模。常见的例子范围从不变量,如原子模拟中的势能、流体动力学中的局部压力 [34] 或视觉模型中的语义类标签 [35],到几何张量,如流体速度场 [34] 和表面法线 [36]。将模型视为一个将输入映射到输出的一般光滑函数,物理一致性要求该输出必须在与系统相关的对称群作用下可预测地变换。当等变性被硬编码到架构中时,这些变换规律是精确的 [见图1(b)]:分子的能量在全局旋转下不变,而速度随参考系协变旋转。形式上,如果对所有,函数关于群是等变的,其中是作用于输出空间的群表示。对于标量目标(例如能量),是恒等映射,从而恢复不变性的定义:。通常,如果输出属于维度为的不可约表示(irrep),则群作用由矩阵介导。

图1. 对称感知机器学习模型的结构、群等变条件以及本工作引入的对称性诊断指标概述。(a) 机器学习模型由一个一般光滑函数表示,它预测输入(例如装饰点云)的物理性质(不同秩的张量)。可以是保持对称性(即等变)的或无约束模型。(b) 当且仅当模型预测在适当的群对称操作作用下像输入一样变换时,群等变性才得以保持。(c) 本工作引入的指标和量化模型预测的等变误差和内部特征的群对称性内容。对于通过Haar积分在群上给出的一组输入,等变误差由反向变换模型预测的标准差给出,而特征标投影给出来自任意层的模型特征的群对称性分解。

当等变性是从数据中学习而非强制时,它本质上是近似的。为了量化模型对表示的变换规则的保真度,我们使用等变误差: (1)

这里,尖括号表示对于任何连续群,对输入的变换在群上的Haar平均,或对于离散群的群平均。注意,项表示输出变换回参考系的量;对于完美等变函数,该量是常数 [等于]。因此,有效地测量了群轨道上重投影输出的标准差 [图1(c)]。该度量已在参考文献 [37] 中基于精心设计的群元素子集引入,后来在参考文献 [22] 中用于估计 Point-Edge Transformer (PET) 机器学习原子间势模型预测的力的等变误差,并在参考文献 [38] 中用于量化使用强制等变的惩罚项训练的无约束模型的近似等变性。参考文献 [24] 使用相同的群平均方差来测量验证误差中有多少来自不完全的输出等变性,并在训练期间跟踪该量。在这里,我们将这种输出级诊断与下面引入的特征标投影配对,以解析中间表示的不可约对称性内容。在补充材料 [39] 中,我们证明了这是一个恰当的等变性度量,因为它当且仅当关于精确等变时才消失。实际上,如补充材料 [39] 所示,该度量可以写成 (2)

这可以更有效地计算,因为它只需要一次群平均,避免了模型在群元素组合上的评估。

然而,输出等变性只是学习过程的终点。两个具有可比值的模型可能依赖于非常不同的内部对称性谱,这些差异可能决定给定架构能否有效学习具有不同角阶或宇称的目标。虽然最终输出被优化为等变,但无约束模型的内部特征不受任何特定对称性约束。隐藏层可能混合不同等变特征的分量,同时仍能有效传播信息。为了探测潜在空间内的对称性内容,我们引入了特征范数的谱分解。类似于信号的功率谱,我们将任何内部特征(例如模型深处任何隐藏层的输出;见图1(c))的群平均范数分解为不同不可约表示的贡献: (3)

本质上,我们寻找一个量,表示变换为不可约表示(更正式地说,是与不可约表示相关的特征标滤波分量的平方范数,与 Peter-Weyl 分解一致 [42])。这里感兴趣的量是特征标投影: (4)

其中是群特征标 [43]。实际上是以群卷积定义的 [44–46]。为方便起见,我们将使用归一化的特征标投影,以简化架构不同特征和不同层之间的比较。关于该量的细节可以在补充材料 [39] 中找到,包括最小化模型函数调用次数的高效评估方案。我们指出,其他方法 [47] 可以为这里引入的度量提供补充信息 [48]。

和都独立于输入参考系的选择(见补充材料 [39]),这使得这两个度量成为模型和输入的内在属性,而不是特定方向选择的伪影。总之,它们提供了关于模型对称性的互补信息,其中适合衡量模型输出遵循所需对称性的准确性,而能够评估污染输出的表示特征,更广泛地识别模型内部特征的谱内容。

虽然在这项工作中我们关注,但和的构造可以自然地扩展到任何紧致群,包括有限群,只需将群平均、不可约表示和特征标替换为感兴趣的对称群的那些。对于紧致群(如),这种推广因此是形式上的且直接的。相反,对于非紧致群(如洛伦兹群),在整个群上归一化的Haar平均不可用,因此当前诊断的实际对应物将需要在群轨道上使用受限或任务相关的采样测度。

III. 方法无约束原子模拟模型

针对原子尺度结构的量子力学计算输出进行拟合的替代模型,是展示基于所引入度量进行对称性分析之概念与启发价值的理想起点。此类模型的输入通常是原子位置与化学物种,形式上可描述为装饰点云。点云边在任意全局旋转、反演和平移的组合下等变,因此此处相关的群为欧几里得群E(3)。平移不变性易于通过使用原子间距离而非绝对位置来强制实现,这也是基于图的模型的典型做法,因此我们可仅关注正交群O(3)。

这些模型所瞄准的物理性质在O(3)的(选定)不可约表示下自然变换,因此无约束模型必须通过数据增强学习适当的变换规律。我们以\(\ell\)表示O(3)群的不可约表示,其中\(\ell\)为角阶,\(p\)为反演下的宇称,其区分了真子空间\(V_\ell\)与其伪补\(V_\ell^-\)。例如,标量是真不变量\((0,0)\),矢量属于\((1,0)\),而赝矢量属于\((1,1)\)。在补充材料[39]中,我们提供了该群等变误差显式形式及特征投影的详细讨论。实际上,所需的群平均使用乘积积分网格(即球面上的Lebedev–Laikov网格[49],加上其余欧拉角上的梯形网格[43,50])进行计算,这为前文引入的度量提供了计算高效且精确的求值(详见补充材料[39]中的证明与细节)。

我们聚焦于PET架构,这是一种基于Transformer的图神经网络(GNN),其输入包括原子点云的边距离和边向量[13]。它在标准基准[13,28]和材料科学应用[23,51,52]中作为机器学习势(MLIP)表现出卓越性能,并且是该领域中日益壮大的基于Transformer架构家族的代表性范例[26,53]。训练中的数据增强有助于模型学习势能面(PES)的O(3)不变性,使得等变误差相对于基线精度可忽略,并在凝聚相中产生稳定的动力学[22]。

我们按复杂度递增的步骤研究PET架构,以凸显诊断与设计的不同方面。首先在第三节A中,我们分析一个在MAD-1.5数据集[54]上训练的通用MLIP,其目标为势能[标量,\((0,0)\)不可约表示]、原子力[矢量,\((1,0)\)不可约表示]和晶胞应力[对称二阶笛卡尔张量,\((0,0)\)与\((2,0)\)不可约表示]。然后,在第三节B中,我们追踪角信息在架构各层中的谱分解,从随机初始化直至完全训练后的模型。凭借此处获得的洞见,在第三节C中,我们提出一种简单的读出权重后处理,进一步纯化输出的对称性内容,从而降低等变误差。最后,在第三节D与E中,我们研究PET内置的归纳偏置如何影响PES之外物理目标的学习动力学,并展示对称性分析如何为简单而有效的架构修改提供依据,以改善学习动力学和模型精度。

A. 学习势能面

我们在大规模原子多样性(MAD)数据集[55]的1.5版本[54]上训练了一个PET MLIP。模型从头训练2000个轮次,以便恰当地考察其随机初始化时的对称性内容并追踪其在训练中的演化。我们首先在MAD-1.5测试集的一个子集上评估训练后模型的能量、力和应力。图2上方面板展示了等变误差分布与绝对模型误差分布的比较。在所有情形下,等变误差分布均向小于绝对误差(AE)的方向偏移,其中能量、非保守力和非保守应力的中位数比值分别为10%、31%和26%。我们关注网络直接预测的力与应力,以研究PET对高阶不可约表示的对称性行为,因为当它们作为能量导数计算时,其张量性质源于微分算子,因此直接与能量的不变行为相关联。

图2. PET MLIP的等变性诊断。上方面板:能量\(E\)、非保守力\(\mathbf{f}\)和非保守应力\(\boldsymbol{\sigma}\)(按\(\ell\)通道拆分)的AE与等变误差\(\mathcal{E}_{\text{eq}}\)的分布。\(\mathcal{E}_{\text{eq}}\)轴上的箭头指示分布中位数。下方面板:相应量的归一化特征投影\(\mathcal{P}_\ell\)作为探测角动量通道\(\ell\)的函数。实线与标记为500个随机采样测试结构的平均值,浅色线为单个结构的投影。

鉴于输出并非精确等变,我们可利用归一化特征投影来探究误差的性质[图2下方面板,展示500个随机采样测试结构的\(\mathcal{P}_\ell\)]。如预期,标量通道主导能量预测,来自更高角动量或伪子空间通道的贡献可忽略[56]。力表现为激活\((1,0)\)通道的矢量,而应力张量(二阶对称笛卡尔张量)激活\((0,0)\)与\((2,0)\)真通道[57,58]。第二点重要观察是伪通道始终比真通道活性低。这一点在能量情形尤为显著,而力与应力的差异则较小。特别是,赝标量\((0,1)\)通道在所有情形下携带的信号强度非常低,且不遵循伪子空间\(V_\ell^-\)中更高角阶项所呈现的指数衰减趋势。

通过监测训练过程中特征投影的演化可获得更多洞见。为此,我们引入特征投影热图[图3(a)],以可视化不同角动量与宇称通道的\(\mathcal{P}_\ell\)演化,从未训练的随机初始化模型一直到完全训练模型。深色表示较小分量,亮色表示较大分量。在任意垂直切片(即特定轮次)处,热图呈现的谱分解与图2下方面板类似。

图3. (a) 特征投影热图报告某量在相关群特征\(\ell\)下以及训练轮次连续过程中的特征投影\(\mathcal{P}_\ell\)大小。此处特征为O(3)的\(\ell\)不可约表示。(b) 能量\(E\)、非保守力\(\mathbf{f}\)及非保守应力两个不可约球面分量\(\boldsymbol{\sigma}_{\ell,p}\)的训练曲线与特征热图。训练曲线中,虚线表示未训练模型的初始均方根误差(RMSE,黑色)或等变误差(蓝色)。

对于PET模型在训练中的每个输出,我们监测测试集模型误差、等变误差及特征投影热图[图3(b)]。可以看到,随机初始化模型对标量\((0,0)\)特征有强烈偏置,而对\((1,0)\)和/或\((2,0)\)的谱含量接近零。随后我们观察到突然的转变,这在非保守力和\((2,0)\)非保守应力预测中最为明显:标量特征下降,\((1,0)\)和\((2,0)\)特征分别占据主导,同时伴随着高\(\ell\)和伪分量(\(\ell^-\))的突然跃升。在后续轮次中,随着目标RMSE和等变误差的减小,特征分解更尖锐地集中于所需不可约表示,除了在轮次约200时出现的\((2,1)\)应力分量外,没有其他突变。标量应力分量是唯一在训练中等变误差相对于初始值增大的目标:原因在于训练开始时标量应力预测接近零,这平凡地不变。非保守应力学习的晚发意味着模型在训练结束时尚未完全收敛,其后果将在第三节C中显现。

B. 窥视黑箱

PET的特征分解分析表明,对称破缺误差中存在大量结构,且模型在不变目标上表现尤佳。为理解原因并提升PET学习非不变目标的能力,我们必须更深入考察其如何学习对称性。为此,我们在上述两个轴向上探索PET特征的信息内容:贯穿架构层以及训练全过程。

图4给出PET MLIP架构的高层概览。对于上一节训练模型中若干重要隐藏层的输出,我们进行了特征分解分析,所得热图标注在架构上。分子或材料中原子中心的环境作为装饰点云输入到架构[图4(a)]。这些环境的内部表示由骨干PET架构生成。该架构由一系列消息传递GNN层组成,每层包含若干模块[图4(b)]。所得骨干特征(BBF)包含所有模型输出所共有的丰富几何信息。这些特征经目标特定的非线性头转换为末层特征(LLF),再由线性读出层映射为模型输出。在最粗的层次上,骨干架构构建的内部表示在训练过程中经历丰富的动力学演化,发展出高\(\ell\)以及赝张量(\(\ell^-\))特征,尤其以边特征为强。

图4. PET MLIP架构概览与内部特征的动力学演化。(a) 分子/材料中原子的局域原子环境表示为装饰点云。中心原子(\(z_i\))及其邻居(\(z_j\))的原子物种,连同边向量(\(\mathbf{r}_{ij}\))及其模长(\(r_{ij}\)),构成PET的输入。(b) 骨干模块:嵌入模块通过中心、邻居和几何嵌入器将输入映射到潜在空间。每个GNN层一个边嵌入器,随后进入Transformer层。(c) 完整架构:预测能量\(E\)、非保守力\(\mathbf{f}\)与非保守应力\(\boldsymbol{\sigma}\)的完整模型流程。架构超参数选取如下:GNN层数\(L_{\text{GNN}}=4\),Transformer层数\(L_{\text{TF}}=3\),截断半径\(r_c=5\) ?,边特征尺寸\(F_e=32\),节点特征尺寸\(F_n=32\),末层特征尺寸\(F_{\ell}=32\)。周围标注红色字母(A–I)的热图对应架构中的特定位置。它们显示归一化特征投影\(\mathcal{P}_\ell\)的强度随训练轮次(\(x\)轴,对数尺度)和探测\(\ell\)通道(\(y\)轴)的变化。在每个热图组内,上下两个面板分别表示\(+\)和\(-\)宇称通道(如热图I右侧所详述)。每个热图最左侧的细长孤立列表示未训练模型,在\(x\)轴上标为“U”。

在分析骨干架构的细节时,我们可以看到先前在模型输出上观察到的特征趋势是如何发展的。对于批次中的给定局域原子环境,中心原子及其邻居的化学性质被嵌入为可学习标量令牌[中心嵌入(CE)和邻居嵌入(NE)],几何嵌入(GE)表示局域三维结构——具体为到每个邻居的距离\(r_{ij}\)和原子间向量\(\mathbf{r}_{ij}\)[图4(b)]。这些GE具有纯粹的标量和矢量性质,在训练过程中不变(热图A)。一个多层感知机将GE转换为固定尺寸的边嵌入(EE)。非线性生成了更高\(\ell\)特征,但仍缺乏赝张量特征(热图B)。

Transformer块,作为PET架构的核心,凭借注意力机制能够生成丰富的几何特征,该机制混合来自多条边的信息并填充赝张量通道(热图C–F)。这对于边特征尤其明显,尽管信号仍较弱。经过两个Transformer层后,第一个GNN层完成,此时模型已能在单个邻域内组合几何信息。即使在随机权重的初始化阶段,GNN块的输出也以标量特征为主,但包含显著的高\(\ell\)分量(热图E与F)。消息传递步骤将几何信息与前一步骤的边消息结合。因此,在第一个GNN层中获得的伪张量激活会传播到下一层,从而在新的EE(热力图G)中可见,而第二个GNN层的输出(热力图H)相对于第一个GNN层的输出没有其他显著变化。最终的BBF是在第二个GNN层的消息传递块之后获得的。边BBF通常携带比节点BBF丰富得多的角向信息。然后,BBF被分配到几个目标特定的非线性头中,并映射到LLF(热力图I)。这些是网络中最后的非线性操作。在随机初始化的模型中,LLF仅具有低阶特征,但在后续阶段,它们会发展出更丰富的几何结构,并一直持续到训练结束。这一点尤其可以在每个目标的LLF的边上再次观察到。在不同目标之间,LLF特征的对称性内容差异显著:例如,与力和应力相比,能量LLF对高阶角向特征的抑制更强。最后,LLF通过线性读出层映射到预测结果。这些读出层只能调节现有分量的幅度,并充当确定预测的最终特征和等变性误差的滤波器。边预测在相邻原子上进行聚合,如果输出是每个结构的量(能量或应力),则节点和边预测都在中心原子上进行聚合。通过将节点和边输出相加,得到如前文图3所示的最终预测结果。这一详细分析为PET如何学习(近似)等变性提供了令人信服的图景。它还为设计严格等变模型提供了有用的指示,因为无约束架构学习到的特征分解表明,需要哪些类型的对称通道才能高效地学习能量、力和应力。它表明,在我们使用的GE下,PET强烈偏向于低 \( L \) 分量,并且即使网络具有生成高阶特征分量的表达能力,它也仅适度使用这种能力,收敛到以低 \( L \) 特征为主的中间表示。这与等变网络通常可以通过将内部表示限制在 \( L \) [59] 来获得相当准确的预测的观察一致。然而,边表示中存在高阶分量表明,高分辨率对于处理几何消息传递信息很重要,因此尝试在至少某些GNN层中包含更高 \( L \) 项并投入部分计算预算的等变模型可能是有用的。 C. 读出层的对称性纯化 在上一节中,我们观察到最终的读出层必须从LLF中滤除不期望的对称性内容。这是通过在训练过程中使用随机旋转进行对称性增强来实现的,从而得到一个等变性误差仅占模型误差一小部分的模型。尽管如此,LLF标记仍被描述目标子空间不需要的不可约表示所污染。我们在这里提出一个简单的协议来纯化线性读出,它可以表述为一个正则化回归问题。如果 \( f_L \) 是输入 \( V \) 的LLF,线性读出可以表示为 \( f_{\text{read}} = W f_L \)。通过计算LLF在群轨道上(或对于连续群如 \( SO(3) \) 则使用网格),我们可以定义两个损失项, (5) 第一个项描述了模型在群上平均的均方误差,第二个项对应于平方等变性误差 \( \mathcal{E}^2_{\text{eq}} \)。组合损失 \( \mathcal{L} = \mathcal{L}_{\text{MSE}} + \lambda \mathcal{L}_{\text{eq}} \),其中 \( \lambda \) 控制等变性惩罚的相对权重,可以转化为一个显式的最小二乘问题,并针对读出权重 \( W \) 求解。确定 \( W \) 所需的所有项都可以通过对训练集进行一次扫描来计算。最小化该损失的最优权重的闭式推导见补充材料 [39]。对于PET模型,我们发现该过程对大多数输出产生了适度的等变性误差改进,证实了传统模型训练期间执行的即时增强成功最小化了 \( \mathcal{E}^2_{\text{eq}} \)(图5)。例外是非保守应力的标量分量。在这种情况下,相对于基础模型,RMSE和等变性误差 \( \mathcal{E}^2_{\text{eq}} \) 都显著改善,分别提高了约5%和超过10%,最大的改进出现在不同的 \( \lambda \) 值处。我们假设这表明训练未完全优化该目标的读出。对于其他目标,我们可以观察到,通过调整 \( \lambda \) 的权重,等变性误差可以降低多达5%(对于对称应力甚至更多),而相对于目标的RMSE牺牲不到1%。在主要训练期间包含显式的等变性惩罚,如参考文献 [38] 中所述,也会影响骨干权重的优化,但需要对每个结构进行多次模型评估,增加了训练的计算成本。相反,我们在此提出的等变读出优化成本等同于对数据集中的每个结构运行一次对称化推理,并且可以作为常规方法验证并可能改进无约束模型的对称性属性。图5. 模型RMSE和等变性误差(\( \mathcal{E}^2_{\text{eq}} \))对于能量 \( E \)、非保守力 \( F \) 和非保守应力 \( \Sigma \) 的分量,针对一个通用PET模型(在MAD-1.5上训练和测试),其能量、非保守力和非保守应力读出层使用结合模型误差和等变性误差惩罚的损失 \( \mathcal{L} \) 进行重新训练。标记颜色对应于等变性惩罚的权重 \( \lambda \)。相对于标准模型的RMSE和 \( \mathcal{E}^2_{\text{eq}} \) 的百分比值显示在图例中。 D. 对模型几何表达能力的压力测试 内部表示以低阶特征为主这一事实引发了以下问题:PET能否学习其对称性并非由初始随机权重自然表达的目标,我们接下来将对此进行研究。模型特征的对称性分析中最清晰的信号是 \( L=2 \) 不可约表示,特别是伪标量,在整个PET架构中具有可忽略的特征。为了确定这是否影响模型的表示能力,我们将一个纯几何伪标量目标 \( \mathcal{P} \) 定义为分子系统中所有键三元组的三重积之和: (6) 并生成一个包含1000个摇动构象的溴氯氟甲烷(CHBrClF)玩具数据集,这些构象仅在初始方向和小的随机坐标畸变上有所不同。我们使用80:10:10的数据划分从头训练PET。图6的上方面板显示了训练轮次中内部特征(EE、BBF、LLF)和预测(\( \mathcal{P} \))的特征投影热力图,而底部面板报告了RMSE和等变性误差(\( \mathcal{E}^2_{\text{eq}} \))。网络在训练过程中表现出显著的两阶段学习行为,类似于在其他模型中经常观察到的所谓的顿悟现象 [60]——尽管在这种情况下,训练集和验证集均如此。从这个角度来看,这一观察表明,网络从简单向量描述符生成高阶对称项的能力可以解释为Transformer参数空间中的相变 [61]。在初始阶段(直到红色虚线,约30个epoch),只有真(\( L=0,2 \))通道活跃,使得伪特征几乎处于休眠状态。因为真(特别是纯不变)通道缺乏表达 \( \mathcal{P} \) 所需的对称性,所以测试RMSE和 \( \mathcal{E}^2_{\text{eq}} \) 都停滞不前。在红色和绿色虚线之间发生了剧烈的转变,伪张量通道突然在BBF、LLF和输出(\( \mathcal{P} \))中激活。这种内部转变引发了RMSE和 \( \mathcal{E}^2_{\text{eq}} \) 的急剧持续下降。正如预期的那样,EE特征没有显示出伪张量激活,因为在这个阶段——恰好在Transformer之前——模型单独处理边,尚不能将它们组合起来。构造伪张量需要混合至少两个向量(“二阶”效应),而形成伪标量需要组合至少三个(“三阶”效应;详见补充材料 [39])。网络必须在注意力层内从严格真输入合成这些伪表示,导致学习的延迟开始。因为伪通道最初很弱,它们的梯度信号被抑制,它们的构造需要跨多个特征的协调高阶交互,使得优化本质上缓慢,直到达到这些分量可以自增强的阈值。鉴于 \( \mathcal{P} \) 是一个纯几何目标,有些琐碎地从分子几何定义而来,但即使在这种情况下,学习任务也如此困难,这值得注意。图6. 几何伪标量训练的等变性诊断。上方:PET模型在训练过程中跨架构层的特征投影热力图。下方:测试集RMSE和平均等变性误差。红色和绿色虚线之间的区域标志着快速学习的开始。这一阶段由伪标量通道在所有能够表达它的层中的强烈激活驱动,并伴有高阶张量(\( L=2 \))和伪张量(\( L=2 \))通道的更广泛、更弱的激活(如右侧所示)。这一诊断揭示了架构设计如何限制模型构建具有特定群对称性的特征的能力,以及这如何影响学习动态。伪子空间的代表性不足对人工任务之外也有影响,可能影响现实世界的原子应用,例如学习NMR化学屏蔽张量 [32](需要伪向量不可约表示)、圆二色性 [62],或针对包含许多具有 \( L=2 \) 特征的分量的哈密顿量和密度矩阵 [63–66]。 E. 增加PET的几何表达能力 接下来,我们挑战PET模型学习一个高 \( L \) 目标。为此,我们从实空间电子密度 \( \rho(r) \)(定义在一组网格点 \( r_g \) 上)开始,针对自旋单态、电荷中性、有机(涵盖11种元素)的QCML分子数据集子集 [67],使用FHI-aims [68] 计算。然后将密度投影到以原子为中心的辅助基集(球谐函数和径向函数的乘积)\( \{ \phi_{nlm}^A \} \) 上,得到电子密度投影 \( \rho_{nlm}^A \): (7) 其中 \( \nu \) 是一个复合索引,标记原子 \( A \)、径向通道 \( n \)、角向阶 \( L \) 和分量 \( m \)。辅助基集选择得足够表达性强,可以以低基集展开误差 \( \epsilon \) 展开电子密度,即 \( \rho(r) = \sum_{\nu} \rho_\nu \phi_\nu(r) + \epsilon \),其中 \( \rho_\nu \) 是展开系数,\( S_{\mu\nu} = \langle \phi_\mu | \phi_\nu \rangle \) 是辅助基重叠矩阵。电子密度投影构成了一个具有挑战性的目标,源自原子建模中核心的量子力学可观测量,并且至关重要的是,它分解为一系列真张量不可约表示,最高到某个最大角向阶 \( L_{\max} \),在我们的案例中为 \( L_{\max} = 3 \)。通过分离并仅针对最高 \( L \) 通道,我们可以清楚地观察到初始化时缺乏高阶角向特征如何影响现实场景中的学习。图7(a)追踪了架构学习密度投影的 \( L=3 \) 分量的尝试。正如2000个epoch上完全平坦的RMSE曲线所示,标准PET未能学习该目标。因为网络仅嵌入边距离和向量,初始GE严格在 \( L=0 \) 处激活。虽然非线性多层感知机在EE中产生了一些更高 \( L \) 的激活,但这是不够的。由于其特征中缺乏有意义的 \( L=3 \) 信息,模型坍缩到一个平凡解:它预测零以最小化损失,正如等变性误差下降和预测与目标之间的范数比趋于零所证明的那样。图7. 高阶角向阶归纳偏差对学习高阶角向阶目标的影响。标准PET(第一列)和修改后的变体——将标准嵌入替换为高达 \( L_{\max}=4 \) 的实心球谐函数——在电子密度投影的 \( L=3 \) 通道上进行训练。上方:真张量(\( L=0,1,2,3 \))热力图,跟踪训练期间关键网络层中角向阶(\( L \))特征的演变。颜色表示特征投影的归一化强度。下方:每个epoch的验证指标,包括预测RMSE(黑色)、等变性误差(青色)和预测与目标范数比(橙色)。我们建议通过对边几何嵌入进行简单且廉价的修改来解决架构中的这一缺陷。我们将标准边距离和向量替换为边向量的实心球谐函数(SSHs)展开,最大角向阶为 \( L_{\max}=4 \)。SSH天然地同时表示角结构和径向信息,因为每个SSH在构造上带有一个明确定义的角动量通道,并且可以以可忽略的开销以自微分方式求值[69]。实际上,由于SSH的函数形式,标准PET的GE在信息内容上几乎等价于一个在$\ell_{\max}$处截断的(嵌入)SSH展开。通过提高该截断值,我们为模型提供了一个远为优越的起点来描述高$\ell$量。为研究这种修改的归纳偏置的影响,我们从零开始训练与图7(a)中相同的PET架构,仅修改GE层以使用递增最大角序$\ell_{\max}$的SSH嵌入。如图7(b)所示,显式增加输入的角信息极大地改变了学习动力学。虽然$\ell_{\max}=1$的表现类似于标准PET且无法学习,但以$\ell_{\max}=4$和8初始化的模型则成功。对于这些模型,$\ell$特征从一开始就在内部表示中可见,RMSE稳定下降,范数比稳定在1附近,证明模型真正在学习目标而非坍缩为0。到训练结束时,内部特征具有丰富的高阶角信息,成功适应了架构的隐式归纳偏置和物理目标的几何性质。

IV. 超越原子模拟

尽管我们将分析重点放在无约束ML模型应用于原子模拟上,但物理学中有许多领域依赖于不严格遵循所有相关对称性的模型。作为一个简短但有启发性的例子,我们考虑液态氩时间投影室(LArTPCs)中微子探测实验中粒子轨迹的分类问题。我们研究基于点的液态氩掩码自编码器(PoLAr-MAE)方法[19]的对称性,该方法依赖于类似于通用PointNet[20]的无约束架构,并基于公开可用的PILArNet数据集[70]构建。粒子轨迹表示为点云,每个点对应一个探测器相互作用,并由探测到的粒子能量修饰。基于此输入,模型将沿粒子轨迹的点分类为四类:电子簇射、径迹、米歇尔电子、$\gamma$射线(参见图8)。PoLAr-MAE总体表现出色的分类精度,除了少数细粒度结构。我们观察到,尽管分类通常对参考系的刚性旋转不变(模型使用随机旋转增强训练,类似于PET),但仍有少数径迹段的类别依赖于取向,特别是对于分类本身不可靠的子标记结构。

图8. PoLAr-MAE架构的示意图,给出了内部标记的特征分解和分割输出,针对1000个分析事件绘制。为表示目的,我们展示了每个事件各点聚合后的分解。一个代表性事件(与图9相同)的模型类概率的等变误差显示,$\gamma$的值与分支点或轨迹段相关,这些轨迹段的分类对参考系的旋转或反射不稳定。

图9. PoLAr-MAE模型的一个代表性输入和输出。对于给定的输入事件(左图),例如液态氩探测器中由粒子能量修饰的粒子径迹簇(点云),模型对离散事件类别进行逐点分类(右图)。红圈包围的点在至少一次采样旋转或反射下改变分类。坐标轴标签以体素索引表示。

在图8中,我们展示了模型架构的总结和内部表示的深入诊断。图中还显示了一个代表性事件(与图9所示相同)的逐点类概率等变误差。该事件的内部描述符的逐点特征分解见补充材料[39]。可以清楚地看到,分类的不确定性反映在PoLAr-MAE预测的四个通道的等变误差中。从分割输出的特征分解(热力图F)中已经可以看出大量非标量行为的存在,具有主导的标量性质和指数衰减的$\ell$特征。赝特征都小得多,与PET的情况类似,此处的赝标量贡献极小。

观察内部层,我们看到位置嵌入A主要具有矢量特征(即位置嵌入器不生成高阶特征),而局部组标记(由微型PointNet单元处理)显示出更丰富的结构,具有显著的$\ell$分量(热力图B)。Transformer编码器块(热力图C-E)没有显示出显著趋势:主导特征为标量,后续块逐渐抑制更高$\ell$特征。然而,卷积分割器增加了非标量特征的污染。总体而言,PoLAr-MAE显示出与PET中节点特征广泛相似的趋势,但缺少PET边特征更丰富的高$\ell$特征,这表明引入显式边表示可能增加模型的表达能力。大的等变误差与分类的不稳定性相关。虽然引入推理时旋转增强可以降低误差,但更优雅且计算高效的方法是将第三节中引入的读出纯化协议应用于最终卷积分割器。通过显式惩罚分类头在$O(3)$群上的等变误差,可以在不增加推理时开销的情况下强制这些具有挑战性的子标记结构实现旋转不变性。

V. 结论

许多应用领域,从图像识别到自然语言处理,都已经接受了这个苦涩的教训:将领域知识纳入机器学习模型的设计,从长远来看,并不总是能胜过为简单、快速且能扩展至大数据集而设计的架构。在相当长一段时间内,物理学似乎是个例外,大量努力以出色的结果投入到对称适应模型的设计中,这些模型被设计为对对称操作满足不变或协变行为。然而,无约束架构近期的成功将焦点从通过构造强制所有对称性,转移到理解哪些对称结构可以被可靠地学习,哪些应通过架构、目标或针对性描述符来施加。由于对称性所扮演的基础性角色,这引发了一个问题:没有这些内置约束的模型如何以及以何种精度学习对称性。我们开发了一个分析框架,使得能够追踪输出特征以及模型内部表示在架构中和训练过程中的变化,该框架跨领域有用,正如我们通过将其应用于原子尺度模拟和LArTPC实验中粒子轨迹分类的旋转对称性所展示的那样。在这两种情况下,ML模型都能以高精度学习旋转对称性,对称性误差远小于模型误差。事实上,两类误差的中位数比率通常在10%到30%之间。

我们提出了一种廉价的模型线性读出纯化步骤。我们发现,所考虑的基于PET的原子模型使用随机增强已经训练到优秀水平,而这种简单后处理步骤带来的改进是适度的,除了应力的标量分量,其等变误差可以大幅降低。模型内部表示的群特征在训练过程中表现出高度非平凡的动力学。最初,低角动量$\ell$特征在网络中占主导,我们部分将其归因于初始几何编码以标量和矢量形式进行。更高$\ell$项很快出现在内部表示中,特别是与原子间边相关的标记,而只有到训练后期,目标属性才达到对所需对称性的精确近似。应力,作为对称笛卡尔张量含有$\ell=2$分量,仅在此后期阶段被学习。这一观察促使我们研究模型在学习赝标量和高$\ell$目标属性时的性能,这些属性与网络中表示不佳的特征相关。我们发现,标准PET架构确实难以学习这些对称性——我们通过注入更高$\ell$描述符到初始结构编码层面来解决这一缺陷,证明了我们的对称性分析的启发式力量。

粒子轨迹分类模型的分析大致相似。最终特征显示出高度的旋转不变性,等变误差与分类的不稳定性强烈相关——这表明将模型训练到高对称水平的重要性。内部特征在整个模型中表现出主要标量特征,至少在执行基于不变输出的任务时似乎是有益的。

像我们引入的这种特征分析有助于在不同架构和架构超参数之间进行选择,并诊断学习对称行为的困难。我们提取的一些见解可能对设计等变模型也相关。PET架构中学习表示特征的快速衰减验证了在等变MLIP内部表示中将角动量阶限制为$\ell=2$或3的常见选择。然而,当使用相同架构处理高$\ell$目标时,仅增加头部的阶数可能不够,因为我们观察到网络早期缺乏某些分量会负面影响相应对称性输出的准确性,尽管我们使用的非线性头原则上可以生成任意阶项。

通用、灵活、可扩展的架构与包含精确物理先验的架构之间的紧张关系不太可能在短期内解决。与此同时,本文引入的分析框架提供了一种严格的方法来量化ML模型的对称行为,产生可操作的见解,从而提高物理科学中无约束架构的准确性和训练动力学。正如我们通过学习基态电子密度的架构修改所展示的,理解模型的失败模式使研究人员能够通过注入最小所需的归纳偏置来实现稳定性和准确性,而不牺牲无约束架构的表达能力和可扩展性。
相关新闻
生物通微信公众号
微信
新浪微博
  • 搜索
  • 国际
  • 国内
  • 人物
  • 产业
  • 热点
  • 科普

热点排行

    今日动态 | 人才市场 | 新技术专栏 | 中国科学人 | 云展台 | BioHot | 云讲堂直播 | 会展中心 | 特价专栏 | 技术快讯 | 免费试用

    版权所有 生物通

    Copyright© eBiotrade.com, All Rights Reserved

    联系信箱:

    粤ICP备09063491号