用几何量和物理量改进E(3)等变消息传递
本文为大家介绍一篇ICLR 2022 under review的论文“GEOMETRIC AND PHYSICAL QUANTITIES IMPROVE E(3) EQUIVARIANT MESSAGE PASSING”。在很多计算物理学和化学任务中考虑协变信息（例如，位置、力、速度和旋转）是非常重要的。作者引入Steerable E(3) Equivariant Graph Neural Networks(SEGNNs)对图网络进行扩展，使得边和节点的属性不局限于不变标量，而且可以包含协变信息，例如向量和张量。SEGNNs由可控多层感知机组成，能够将几何和物理信息加入到消息公式与更新公式中。作者进一步提出SEGNNs的两个成功组件：非线性消息聚合改进了传统的线性（可控）点卷积；可控消息改进了近年来发送不变消息的等变图网络。最后，作者在几项计算物理学和化学任务上进行实验并提供额外的消融研究，证明了他们方法的有效性。
1介绍
卷积神经网络CNNs的问世是深度学习迅速崛起的重要因素，归因于CNNs开发平移对称性的强大能力。最近工作表明通过群卷积来开发其它对称性能进一步提高CNNs的性能。图神经网络（GNNs）和CNNs是紧密相关的，因为它们都是对局部信息的聚合。消息传递层的长处在于能以高度非线性的方式转换和传播节点特征。近年来一些等变GNNs的重要组成组件都是卷积层，尽管非常强大，但是这样的层只能对图进行线性转换，而非线性只存在于激活函数中。
在这篇论文中，作者提出非线性E(3)等变消息传递层，它使用了和群卷积一样的底层原理，也可以将其视为非线性群卷积。该方法的中心思想是使用可控向量以及它们的等变转换来表示和处理节点特征。于是，节点和边的特征可以是旋转不变的（标量）或者是协变的（向量，张量）。在可控消息传递框架中，作者使用了Clebsch-Gordan（CG）张量积通过几何信息对消息公式和更新公式进行控制。节点属性可以包括如节点速度、力、原子旋转等信息。目前，尤其是在分子建模领域中，大部分的数据集都是由原子数量和位置信息构建而成。在本文中，作者展示出用更多几何和物理量来丰富节点属性的潜力。
这篇论文的主要贡献有以下四点：
- 对等变图网络进行扩展，使得节点和边属性不再局限于标量。
- 用于可控向量场的一类新的激活函数，基于可控节点属性和可控多层感知机的极少，这类激活函数允许将几何和物理量注入到节点更新中。
- 通过非线性卷积的定义对不同的等变图神经网络有统一的见解。
- 大量的消融研究，展示出可控消息传递比不可控的消息传递效果要好，非线性卷积效果比线性卷积要好。
2泛化的E(3)等变可控消息传递
可控特征
可控多层感知机
球谐编码
使用球谐变化，将任何三维向量转变为type-l向量。如下所示：
作者将使用球谐嵌入将几何和物理信息加入到可控MLPs中。球谐公式的可视化如图2。
可控激活函数
在可控环境下，需要确保激活函数也是等变的。作者将Weiler等人在2018年提出的gated激活函数加入到自己的架构中。最终，该可控MLPs其本身就可以作为直接利用局部几何线索的一类新的激活函数。
3实验结果
SEGNN结构和消融
在所有的任务中，作者都至少设置了一个完全可控的SEGNN。作者进行了消融实验来研究使SEGNN和文献不同的两个主要原则。A1非可控vs可控的情形通过在lf、la中设置不同的最大球谐阶所得到。那么EGNN可以当成lf=0,la=0的特殊情形。这些模型用SEGNN标记。A2使用可控等变点卷积方式（Thomas等人在2018年的工作），将其和通过两层可控MLPs非线性方式实现的网络进行对比。
N-body系统
带电N-body系统实验包含五种带正电荷或者负电荷的例子，并给定它们的初始三维位置和速度。任务是估计在1000个时间步后粒子的位置。
结果和消融研究如表1所示。当lf=1,la=1可控架构取得最佳效果，增加阶数并不会提高性能。SEGNN(G)模型只利用了几何信息，具体是将边的属性（如相对取向）进行平均作为节点属性。而SEGNN(G+P)在前者的基础上，将速度的球谐编码加入到节点属性中。从表1中可以看出，SEGNN(G+P)性能相对于SEGNN(G)又有所提升。
QM9数据集
QM9数据集包含拥有29种原子的小分子，并使用三维位置坐标和表示原子类型的独热编码来描述每一个原子。任务是通过优化预测和基本事实之间的平均绝对误差来回归出各种分子的化学性质。
如表3所示，随着可控特征向量阶数提高，性能也随之提高，对于小的cutoff radius尤其明显。表2和表3共同表明SEGNNs表现优于架构可比的EGNN模型，同时从中剥离出注意力模块并且将图的连接性进行缩减。
OC20数据集
Open Catalyst Project OC20数据集由表面上的molecular adsorptions组成。作者关注的是Initial Structure to Relaxed Energy（IS2RE）任务，该任务将输入作为初始结构并预测最终的能量。最终的测试集分为四类，in-distribution（ID）催化剂和吸附物、out-of-domain吸附物（OOD Ads）、out-of-distribution催化剂(OOD Cat)、out-of-distribution（OOD Both）吸附物和催化剂。
表4展示了在OC20数据集上SEGNN模型和其它模型的对比。通过消融研究得到了最好的SEGNN模型，该模型基本在所有指标上都取得最先进的结果。
4总结
作者介绍了SEGNNs，它在等变图网络上进行扩展使得节点和边的信息不再局限于标量，而可以是张量或向量。SEGNNs是第一个允许几何和物理量对于节点更新进行控制的网络，同时作者还介绍了一类新的等变激活函数。最后作者将SEGNNs运用到大量不同的任务中。广泛的消融研究进一步说明了可控相对于不可控消息传递的好处、非线性相对于线性卷积的好处。
参考资料
https://openreview.net/forum?id=_xwr8gOBeV1