E(3)-equivariant graph neural networks for data-efficient and accurate interatomic potentials

E(3)-equivariant graph neural networks for data-efficient and accurate interatomic potentials

刊物

NATURE COMMUNICATIONS

解决问题

• 大多数当代的对称性感知模型使用不变卷积，并且只作用于标量

创新点

• NequIP使用E(3)-等变卷积来进行几何张量的相互作用，从而得到信息更加丰富更加忠实的原子环境表示
• 达到了SOTA精度
• 高数据效率

思想概述

等变

• 在原子系统中，物理性质在一组原子的平移、反射和旋转下具有明确定义的变换性质
• 例子：如果一个分子在空间中旋转，它的原子偶极子或力的向量也会通过等变变换相应地旋转
• 等变神经网络能够更一般地表示物理系统的张量性质和张量运算（例如向量加法、点积和交叉积）
• 形式上，一个函数$f:X\mapsto Y$对于一个作用于$X$和$Y$的群$G$是等变的如果：

$$D_Y[g]f(x)=f(D_X[g]x)\forall g\in G,\forall x\in X$$

• 其中，$D_X[g]$和$D_Y[g]$是群元素$g$在向量空间$X$和$Y$的表示

神经等变原子间势（NequIP）

• 目标：找到从原子位置 { r ⃗ i } \{\vec{r}_{i}\} {r i​}和化学种类 { Z i } \{Z_{i}\} {Zi​}到总势能$E_{pot}$和作用在原子上的力 { F ⃗ i } \{\vec{F}_{i}\} {F i​}的映射，总势能是原子势能的总和，力为总势能相对于原子位置的梯度

$$\begin{gathered} E_{pot}=\sum _{i\in N_{atoms}}{E}_{i,atomic} \\ {\vec{F}}_i=-{\nabla }_i{E}_{pot} \end{gathered}$$

• $E_{i,atomic}$是由图神经网络预测的标量节点属性，NequIP的输出是预测的势能$E_{pot}$，它在平移，反射和旋转下是不变的，但网络包含的内部特征是与旋转反射等变的几何张量
• 网络建立在张量场网络Tensor-Field Networks(TFN)中引入的层上，其基元是在e3nn中实现的，使得能够构建对于平移表现中不变性，对于奇偶性和旋转表现出等变性的神经网络，原子特征由标量、向量和高阶张量组成，特征向量$V_{acm}^{l,p}$的下标是$l,p$，旋转阶$l=0,1,2,...$是非负整数，并且奇偶性是$p\in (1,-1)$中的一个，它们一起标记了$O(3)$的不可约表示，下标$a,c,m$分别对应于原子，通道（特征向量的元素）和取值$m\in [-l,l]$的表示下标，对这些几何对象操作的卷积是等变函数而不是不变函数，即，如果在$k$层的特征在旋转或奇偶变换下进行变换，那么$k\to k+1$层的卷积输出进行相应的变换，卷积操作天然是平移不变的，因为它们的滤波器作用于相对原子间距离向量，此外卷积还是排列不变的，虽然原子特征对于下标是排列等变的，总势能对于下标是排列不变的
• 为了实现旋转不变性，卷积滤波器$S_m^{(l)}({\vec{r}}_{ij})$被约束为可学习的径向函数和球谐波的乘积，它们在$SO(3)$下是等变的：

$$S_m^{(l)}({\vec{r}}_{ij})=R(r_{ij})Y_m^{(l)}({\hat{r}}_{ij})$$

• 其中，如果${\vec{r}}_{ij}$为中心原子$i$到近邻原子$j$的相对位置，则${\hat{r}}_{ij}$和$r_{ij}$分别是相关的单位向量和原子间距离，并且$S_m^{(l)}({\vec{r}}_{ij})$是相应的卷积核，滤波器中所有可学习权值都在旋转不变径向函数$R(r_{ij})$中，这个径向函数被实现为一个多层感知机，它一起输出所有滤波器-特征张量产生的相互作用的径向权值：

$$R(r_{ij})=W_n\sigma (...\sigma (W_2\sigma (W_{1}B(r_{ij}))))$$

• 其中$B(r_{ij})\in {\mathbb{R}}^{N_b}$是$N_b$维的原子间距离的基嵌入，$W_i$是权重矩阵并且$\sigma (x)$表示逐元素的非线性激活函数，文中使用的是SiLU激活函数，径向贝塞尔函数和一个多项式包络函数用于作为原子间距离的基：

$$B(r_{ij})=\frac{2}{r_c}\frac{\sin (\frac{b\pi }{r_c}{r}_{ij})}{r_{ij}}{f}_{env}(r_{ij},r_c)$$

• 其中$r_c$是一个局部截止半径，将相互作用限制在比某个截止距离更近的原子上，因为径向贝塞尔函数在截止距离$r_c$处不是二阶连续可微的，所以需要乘以一个包络函数，$f_{env}$是一个多项式，$p=6$作用于由截止半径归一化的原子间距离$\frac{r_{ij}}{r_c}$

$$f_{env}(r_{ij},r_c)=1-\frac{(p+1)(p+2)}{2}{(\frac{r_{ij}}{r_c})}^p+p(p+2){(\frac{r_{ij}}{r_c})}^{p+1}-\frac{p(p+1)}{2}{(\frac{r_{ij}}{r_c})}^{p+2}$$

• 截止/局部原子环境的使用允许评估的计算成本与原子数量线性成比例，网络初始化时，贝塞尔根被设置为$b=[1,2,...,N_b]$，其中$N_b$是基函数的数量，随后通过反向传播来优化$b\pi$，对于具有周期边界条件的系统，使用在ASE代码中实现的邻居列表来识别适当的原子邻居
• 最后在卷积中，输入的原子特征张量和滤波器必须再次以等变的方式结合起来，这点通过一个几何张量积达到，其输出特征是旋转等变的
• 网络需要尊重的最后一个对称性是奇偶性：网络在反转下如何变换，即$\vec{x}\to -\vec{x}$，如果满足以下选择规则，奇偶等变性是通过只允许来自一个滤波器和一个具有奇偶值$p_f$和$p_i$的输入张量特征的贡献来贡献一个输出特征来实现的：

$$p_o=p_i{p}_f$$

• 最终，一个完整的卷积层与滤波器$f$作用于输入$i$，实现产生输出$o:l_i\otimes l_f\to l_o$：

$${\mathcal{L}}_{ac{m}_o}^{l_o,p_o,l_f,p_f,l_i,p_i}({\vec{r}}_a,V_{ac{m}_i}^{l_i,p_i})=\sum _{m_f,m_i}{C}_{l_i,m_i,l_f,m_f}^{l_o,m_o}\sum _{b\in S}(R(r_{ab}{)}_{c,l_o,p_o,l_f,p_f,l_i,p_i})Y_{m_f}^{l_f}({\hat{r}}_{ab})V_{bc{m}_i}^{l_i,p_i}$$

• 其中，$a$和$b$分别表示卷积的中心原子和相邻的原子$b\in S$，$C$表示Clebsch-Gordan系数，对于给定的输出旋转阶和奇偶性$(l_o,p_o)$，由$(l_i,p_i)$和$(l_f,p_f)$的不同组合产生多个${\mathcal{L}}_{ac{m}_o}^{l_o,p_o}$张量，取所有可能的输出张量，并将它们链接起来，还将邻域和的输出除以$\sqrt{N}$，$N$表示原子的平均邻居数，为了更新原子特征，还是用dense层，以逐原子的方式应用，原子之间共享权重，类似于SchNet中的自相互作用层，不同的权重用于不同的旋转阶，对给定的张量的所有下标$m$应用相同的权重集，以保持等变
• NequIP网络架构，原子embedding，interaction blocks，an output block：

• Embedding：跟随SchNet，初始特征根据原子序数$Z_i$得到
• Interaction Block：编码相邻原子之间的相互作用：核心是卷积函数，产生相同旋转和奇偶性的对$l_o,p_o$的不同张量积相互作用的特征由线性逐元素self-interaction层混合，使用了resnet风格的更新：${\mathbf{x}}^{\mathbf{k}+1}=f({\mathbf{x}}^{\mathbf{k}})+\text{Self-Interaction}({\mathbf{x}}^{\mathbf{k}})$，其中$f$是一系列self-interaction，convolution，concatenation和self-interaction，上述公式中的每个种类的Self-Interaction的权重是分开学习的，最后，混合特征被等变的基于SiLU的门控非线性处理（偶和奇标量不是门控，而是直接分别由SiLU和tanh非线性处理）
• Output Block：最后卷积的$l=0$特征被传递给一个输出块，该输出块由一组两个逐原子的self-interaction层组成
• 对于每个原子，最后一层输出一个标量，这被解释为原子势能，然后将这些结果求和，给出系统的总预测势能，随后，力作为预测的总势能的负梯度得到，从而确保了能量守恒和旋转等变力
  成
• 对于每个原子，最后一层输出一个标量，这被解释为原子势能，然后将这些结果求和，给出系统的总预测势能，随后，力作为预测的总势能的负梯度得到，从而确保了能量守恒和旋转等变力