等变层神经网络：在图上学习几何传输

# 等变层神经网络：图上的几何传输学习
来源：https://arxiv.org/html/2608.28853  
Alessio Borgi  
所属机构：剑桥大学计算机科学与技术系  
所属机构：罗马第一大学计算机、控制与管理工程系  
Mario Severino  
11脚注标记：1  
所属机构：剑桥大学计算机科学与技术系  
所属机构：帕多瓦大学信息工程系  
Fabrizio Silvestri  
所属机构：罗马第一大学计算机、控制与管理工程系  

###### 摘要  
等变图神经网络为建模几何系统提供了一种有原则的方法，但高效的低阶架构在向量信息沿图移动时的变换能力仍然有限。我们引入了**ESNN**（等变层神经网络），该网络通过学习相邻向量特征之间有向的、矩阵值的传输来丰富这种交互，同时保持精确的欧几里得等变性。ESNN 并未提高表示的阶数，而是将标量和向量特征保持为低阶，并将额外的几何灵活性置于边的传输本身。我们从理论上表征了这种传输，表明当相对位移是唯一的协变几何输入时，每个线性的 $O(n)$ 等变映射都分解为独立的径向和切向分量，而学习到的协变特征则实现了更丰富的、依赖特征的变换。我们还为具有特定环境方向的系统引入了受控的对称性松弛，该方向可以是从数据中预设或推断出来的，而当方向性路径不活跃时，可以恢复完整的 $E(n)$ 等变性。在粒子动力学、基于网格的模拟、点云分类和分子性质预测等任务中，ESNN 改进了动力学预测，在对称性破缺时恢复了重力轴，在选定的网格任务和长期推演中带来了显著提升，并且对未见过的旋转保持稳健。这些结果表明，学习几何信息如何跨边传输，提供了一种补充性的途径来实现富有表达力的等变消息传递，而无需更高阶的表示。  

## 1 引言  
许多物理系统，从分子动力学到流体力学，都自然嵌入在 $n$ 维欧几里得空间中，并受欧几里得群 $E(n)$ 的对称性支配（Satorras 等，2021 (https://arxiv.org/html/2608.28853#bib.bib1)；Brandstetter 等，2022 (https://arxiv.org/html/2608.28853#bib.bib8)；Du 等，2023 (https://arxiv.org/html/2608.28853#bib.bib2)；Wang 等，2024b (https://arxiv.org/html/2608.28853#bib.bib4)；Aykent 和 Xia，2025 (https://arxiv.org/html/2608.28853#bib.bib6)）。当表示为几何图时，其节点状态可以将不变的标量属性与随几何变换的向量或张量特征相结合（Schütt 等，2021 (https://arxiv.org/html/2608.28853#bib.bib7)；Simeon 和 de Fabritiis，2023 (https://arxiv.org/html/2608.28853#bib.bib3)；Liao 等，2024 (https://arxiv.org/html/2608.28853#bib.bib5)）。因此，一个物理上一致的模型应该对平移、旋转以及（在适当情况下）反射做出可预测的响应。等变性约束了这种变换行为，但它本身并不能决定相邻节点之间应如何交换几何信息。在许多物理系统中，这种交互本质上是方向依赖的：纵向和横向响应、剪切和各向异性传播取决于相对方向，而不仅仅是成对距离。因此，几何消息传递应捕获这种方向结构，同时保持精确的 $E(n)$ 等变性。  

EGNN、SNN、ESNN 边消息 $j \to i$  $\mathbf{r}_{ij}$ $\mathbf{r}_{ij}$ $i$ $j$ $\phi(\lVert\mathbf{r}_{ij}\rVert)$ 各向同性标量传输  
示意图空间基元：压缩 $\checkmark$ 剪切 $\times$ 扭转 $\times$  
等变性测试输入 $\mathbf{x}_i$ $f(\mathbf{x})$ $\mathbf{Q}$ 旋转 $\mathbf{Qx}_i$ $\exp$ $\mathbf{Q}f(\mathbf{x})$ $\checkmark$ $f(\mathbf{Qx})=\mathbf{Q}f(\mathbf{x})$ 等变性 $\checkmark$ 各向异性 $\times$  

层消息 $j \to i$  $\mathbf{r}_{ij}$ $\mathcal{F}_{i}$ $\mathcal{F}_{j}$ $= \mathbf{W}_{ij}$ 无约束耦合映射  
示意图空间基元：压缩 $\checkmark$ 剪切 $\checkmark$ 扭转 $\checkmark$  
等变性测试输入 $\mathbf{x}_i$ $f(\mathbf{x})$ $\mathbf{Q}$ 旋转 $\mathbf{Qx}_i$ $\exp$ $\mathbf{Q}f(\mathbf{x})$ $f(\mathbf{Qx})$ $\neq$ $\times$ $f(\mathbf{Qx}) \neq \mathbf{Q}f(\mathbf{x})$ 等变性 $\times$ 各向异性 $\checkmark$  

层消息 $j \to i$  $\mathbf{r}_{ij}$ $\mathcal{F}_{i}$ $\mathcal{F}_{j}$ $\cdot$ $\mathbf{S}_{ij}$ 空间 $\mathbf{M}_{ij}$ 通道  
示意图空间基元：压缩 $\checkmark$ 剪切 $\checkmark$ 扭转 $\checkmark$  
等变性测试输入 $\mathbf{x}_i$ $f(\mathbf{x})$ $\mathbf{Q}$ 旋转 $\mathbf{Qx}_i$ $\exp$ $\mathbf{Q}f(\mathbf{x})$ $\checkmark$ $f(\mathbf{Qx})=\mathbf{Q}f(\mathbf{x})$ 等变性 $\checkmark$ 各向异性 $\checkmark$  

图 1：*EGNN 式消息传递、通用层神经网络和 ESNN 中的逐边几何传输*。EGNN 式模型从不变标量系数和相对位移构建几何交互，而通用层网络允许局部特征空间之间的矩阵值变换，但不一定遵循几何向量的变换法则。ESNN 通过将 $O(n)$ 协变的空间作用与不变通道混合相结合，实现了矩阵值传输和精确等变性。变形图示说明了逐边的空间作用，并非旨在声称每种架构的完整表达能力。  

现有的等变架构通过不同的机制捕获方向结构。笛卡尔模型保持表示空间简单，使用不变标量和低阶向量（Satorras 等，2021 (https://arxiv.org/html/2608.28853#bib.bib1)；Schütt 等，2021 (https://arxiv.org/html/2608.28853#bib.bib7)），而可转向架构则丰富了传播的表示，使用高阶特征、球谐函数和张量积来捕获更详细的角度结构（Thomas 等，2018 (https://arxiv.org/html/2608.28853#bib.bib9)；Batatia 等，2022a (https://arxiv.org/html/2608.28853#bib.bib10)；Batatia 等，2022b (https://arxiv.org/html/2608.28853#bib.bib11)）。一个互补的可能性是保持特征表示为低阶，而是丰富信息在相邻节点之间传递时的变换方式。这正是层神经网络（SNNs）所提出的视角，其中相邻节点之间的交互由局部特征空间之间的学习线性映射来介导（Hansen 和 Gebhart，2020 (https://arxiv.org/html/2608.28853#bib.bib17)；Bodnar 等，2022 (https://arxiv.org/html/2608.28853#bib.bib16)）。这些映射提供了矩阵值的边交互，为更丰富的向量传输提供了一种自然机制，但通用的层参数化并未考虑几何向量特征的欧几里得变换法则。为了弥合这一差距，我们提出了*等变层神经网络（ESNN）*，这是一族等变图神经网络，它在保持精确 $E(n)$ 等变性的同时，学习相邻向量特征之间的矩阵值传输。ESNN 将每次传输分解为一个 $O(n)$ 协变的[11]空间作用和跨特征通道的不变混合。这使得在完全使用低阶标量和向量表示的同时，建模各向异性交互成为可能。一个特别简单的实例区分了向量消息沿相对位移的分量和垂直于它的分量，并分别对两者进行变换。我们证明，当相对位移是唯一的协变几何输入时，这种径向-切向分解表征了线性 $O(n)$ 等变传输的完整类别。ESNN 进一步适用于外部结构（如重力或背景流）选择特定方向并降低动力学对称性的系统。该方向可以从先验知识中固定，也可以作为全局模型参数学习，而可学习的松弛系数控制其对传输影响的强弱。当系数为零时，模型是精确 $E(n)$ 等变的；当使用方向性信号时，对称性降低到保持该特定方向的子群。  

**贡献**。我们的主要贡献包括：  
- • 我们引入了*等变层神经网络（ESNN）*，这是一个低阶笛卡尔框架，它在保持精确 $E(n)$ 等变性的同时，学习向量特征跨边的矩阵值传输。  
- • 我们开发了一系列等变传输机制，具有递增的几何灵活性，从恒等和各向同性变换到依赖特征的旋转和各向异性径向-切向传输。我们证明，当相对位移是唯一的协变几何输入时，径向-切向形式捕获了线性 $O(n)$ 等变传输的完整类别。  
- • 我们为具有特定环境方向的系统引入了受控的对称性松弛，该方向可以从数据中预设或学习。一个可学习的系数控制其影响：当为零时，精确恢复完整的 $E(n)$ 等变性；当活跃时，等变性保持到保持该特定方向的子群。  
- • 我们在粒子动力学、基于网格的物理模拟、点云分类和分子性质预测中评估了 ESNN，测试了在完全对称下的更丰富传输、通过数据驱动恢复对称性破缺方向，以及在不同几何领域间的迁移能力。  

## 2 相关工作与背景  
**相关工作**。ESNN 处于等变消息传递和基于层的图学习的交叉点，结合了前者的几何归纳偏见和后者的矩阵值传输视角。笛卡尔架构如 EGNN（Satorras 等，2021 (https://arxiv.org/html/2608.28853#bib.bib1)）以及标量-向量模型如 PaiNN 和 GVP（Schütt 等，2021 (https://arxiv.org/html/2608.28853#bib.bib7)；Jing 等，2021 (https://arxiv.org/html/2608.28853#bib.bib18)）使用不变标量和低阶协变特征实现了高效的欧几里得等变性。另一条互补的工作线开发了可转向模型，这些模型根据特征在旋转下的变换方式来表示特征，并通过以保持对称性的方式组合这些表示类型来构建等变交互（Thomas 等，2018 (https://arxiv.org/html/2608.28853#bib.bib9)；Fuchs 等，2020 (https://arxiv.org/html/2608.28853#bib.bib19)；Brandstetter 等，2022 (https://arxiv.org/html/2608.28853#bib.bib8)；Batatia 等，2022b (https://arxiv.org/html/2608.28853#bib.bib11)；Liao 等，2024 (https://arxiv.org/html/2608.28853#bib.bib5)）。ESNN 保留了低阶笛卡尔设置，但通过学习相邻向量特征在聚合前如何被变换来增加边交互的表达能力。这一视角自然地连接到 SNN，它将标量边权重泛化为局部特征空间之间的学习线性映射（Hansen 和 Gebhart，2020 (https://arxiv.org/html/2608.28853#bib.bib17)；Bodnar 等，2022 (https://arxiv.org/html/2608.28853#bib.bib16)）。基于连接的层模型考虑了正交几何映射（Barbero 等，2022a (https://arxiv.org/html/2608.28853#bib.bib38)），而最近的工作已将基于层的传播扩展到有向和非对称交互（Ribeiro 等，2025 (https://arxiv.org/html/2608.28853#bib.bib44)；Fiorini 等，2025 (https://arxiv.org/html/2608.28853#bib.bib45)）。协同层（Copresheaf）公式提供了相关视角，其中信息通过局部特征空间之间的有向映射传播（Hajij 等，2025 (https://arxiv.org/html/2608.28853#bib.bib46)）。ESNN 建立在这种局部的、矩阵值的视角之上，同时施加了几何向量特征所需的欧几里得变换法则。ESNN 还允许在环境具有特定方向时放松这种对称性先验。相关的亚等变模型强制执行与指定场相关的子群（Han 等，2022 (https://arxiv.org/html/2608.28853#bib.bib27)），而松弛等变方法则学习对底层对称性的偏离（Wang 等，2022a (https://arxiv.org/html/2608.28853#bib.bib50)；Hofgard 等，2024 (https://arxiv.org/html/2608.28853#bib.bib52)）。在 ESNN 中，这种额外的自由度通过一个可学习的特定方向信号引入，当信号不活跃时，精确恢复完整的 $E(n)$ 等变性。更广泛的讨论见附录 A (https://arxiv.org/html/2608.28853#A1)。  

### 2.1 几何背景  
我们考虑一个由图 $\mathcal{G}=(\mathcal{V},\mathcal{E})$ 表示的物理系统，其中每个节点 $i \in \mathcal{V}$ 具有空间坐标 $\mathbf{x}_{i} \in \mathbb{R}^{n}$。我们的目标是构建尊重欧几里得对称性的消息传递算子，同时允许相邻节点之间的交互依赖于它们的相对几何。我们首先指定相关的变换法则和特征表示，然后引入胞腔层（cellular-sheaf）构造，这为第 3 节 (https://arxiv.org/html/2608.28853#S3) 的传输公式提供了动机。  

**等变性**。欧几里得群 $E(n)=O(n)\ltimes\mathbb{R}^{n}$ 通过刚体运动作用于坐标：$\mathbf{x}_{i} \mapsto Q\mathbf{x}_{i} + \mathbf{t}, \quad Q \in O(n), \quad \mathbf{t} \in \mathbb{R}^{n}$ (1)。一个 $E(n)$ 等变层以与此作用一致的方式变换其输出。图交互通过相对位移 $\mathbf{r}_{ij}=\mathbf{x}_{i}-\mathbf{x}_{j}$ 表达，对于它们，平移相消，而正交变换作用为 $\mathbf{r}_{ij} \mapsto Q\mathbf{r}_{ij}$。对于基于相对几何构建的操作，平移等变性因此内置于表示中，而剩余的几何要求是强制执行适当的 $O(n)$ 变换法则。  

**特征表示**。每个节点携带不变标量特征和协变向量特征：$\mathbf{h}_{i}=(\mathbf{s}_{i},\mathbf{V}_{i})$ (2)，其中 $\mathbf{s}_{i} \in \mathbb{R}^{c_{s}}$ 包含不变标量通道，$\mathbf{V}_{i} \in \mathbb{R}^{n \times c_{v}}$ 包含 $c_{v}$ 个向量通道。$\mathbf{V}_{i}$ 的每一列都是一个 $n$ 维向量，其变换如下：$\mathbf{s}_{i} \mapsto \mathbf{s}_{i}, \quad \mathbf{V}_{i} \mapsto Q\mathbf{V}_{i}$。