张量场神经网络TFN
张量场网络(Tensor Field Networks, TFNs)源于论文Tensor Field Networks: Rotation- and translation-equivariant neural networks for 3D point clouds,是一个为分子和物质建模量身定制的工具,其设计充分考虑了化学和物理学的原理,特别是对称性和不变性。通过这种方式,TFN提供了一个强大的框架,可以精确地捕捉分子的三维结构和它们如何相互作用。
1、TFN输入数据的Embedding
嵌入层的作用是将每个点的输入特征(例如原子的类型和状态)映射到一个高维空间,从而能更丰富地表示其化学性质和物理状态。在分子建模中,通过嵌入层将每个原子的类型和状态映射到一个高维空间,张量场网络能够更丰富地表示分子的化学性质和物理状态。这种嵌入表示为网络的后续层提供了基础,同时也是可学习的,确保了网络在训练过程中能够优化这些特征来更好地捕捉分子属性。特别是在处理三维空间中的分子时,张量场网络进一步将这些嵌入向量与SO(3)群的表示相关联,允许网络通过一系列不可约表示来捕捉不同旋转对称性级别的信息。
具体来说,张量场网络的每一层的输入和输出都是一个有限节点集\(\mathrm{S}\),这些点位于3D空间\(\mathbb{R}^3\)中,每个节点都有一个与SO(3)群的表示相关联的向量。换句话说,网络中的每个节点不仅仅是三维空间中的一个坐标点,它还携带了一定的附加信息,这些信息以向量的形式体现,且这个向量是如SO(3)群元素那样可以变换的。在物理意义上,这意味着这些向量能够反映点的旋转性质和对旋转的响应方式,确保网络在处理旋转时能够保持一致性和不变性。这些向量实际上就是Winger D矩阵(详见第一章)不同的不可约表示,对应于旋转对称性的不同球谐函数阶数,这被称作\(\mathrm{l}\) -旋转阶数不可约表示。例如,\(l=0\)对应于标量(无方向的量),\(l=1\)对应于向量(有方向的量),而更高的\(\mathrm{l}\) 值对应于更复杂的对称性。它们可以被类比于标准卷积神经网络中,卷积操作的卷积核的数量。在CNN中,每个卷积核可以捕捉输入数据的不同特征;类似地,在张量场网络中,不同的\(\mathrm{l}\) -旋转阶数不可约表示可以捕捉数据的不同旋转对称特征。这样,当网络处理3D数据时,它可以同时处理多个旋转对称性级别的信息,从而允许它更全面地理解数据的旋转几何结构。
在张量场网络中,我们构建一个特征张量字典\(V^{(L)}\),用来存储不同旋转阶数\(\mathrm{l}\)下的特征向量。每个键\(\mathrm{l}\)对应一个值,这个值是一个多维数组,形式为\([S,n_{_l},2l+1]\),其中\(\mathrm{S}\)代表点的数量,\(n_{l}\)是提取特定类型特征的“卷积核”数量,而\(\text{2l+1}\)则是参照球谐函数的定义描述的给定\(\mathrm{l}\)下的空间方向维度,论文的名字张量场(Tensor Fields)指代的就是这个张量字典\(V^{(L)}\)。
下面通过一个简化的例子来理解这种特征张量字典\(V^{(L)}\)的概念。假设我们有一个系统,这个系统里有两个星球,我们想要跟踪每个星球的一些特性:质量、位置和速度。此时可以将这些特性编码为特征张量,并使用张量场网络来处理。
\(l=0\):对于张量场中的零阶(标量场),此时只关心每个星球的质量。因此,对于每个星球(或点),可以用一个数值(标量)来表示其质量。令星球1的质量是\(m_{1}\),星球2的质量是\(m_{2}\)。字典的形式可以是:
\(\begin{aligned}&\boldsymbol{V}^{(0)}=\{\\&1{:}[m_1]\text{ //星球1的质量}\\&2{:}[m_2]\text{ //星球2的质量}\\&\}\end{aligned}\)
\(l=1\):对于一阶张量场(也就是向量场),通常涉及三维空间中的向量,因此可以用来表示诸如位置、速度、加速度等物理量。这些向量的每个分量都对应着空间的一个维度(通常是x、y、z坐标轴)。假设当前想要表示每个星球的位置,速度和加速度。令星球1的位置是\((x_1,y_1,z_1)\),速度是\((vx_1,vy_1,vz_1)\),加速度是\((ax_1,ay_1,az_1)\)。星球2的位置是\((x_2,y_2,z_2)\),速度是\((vx_2,vy_2,vz_2)\),加速度是\((ax_2,ay_2,az_2)\)。字典的形式可以是:
$V^{(1)}=\{$
$1: \{$ //星球1
position:$[ x_1, y_1, z_1]$ //位置向量
velocity$:[\boldsymbol{\nu x}_1,\boldsymbol{\nu y}_1,\boldsymbol{\nu z}_1]$ $\|$速度向量
acceleration$:[ax_1,ay_1,az_1]$ //加速度向量$\}$,
$2: \{$ //星球2
position:$[ \boldsymbol{x}_2, \boldsymbol{y}_2, \boldsymbol{z}_2]$ $\|$位置向量
velocity:$[\boldsymbol{\nu x}_2,\boldsymbol{\nu y}_2,\boldsymbol{\nu z}_2]$ //速度向量
acceleration:$[ax_2,ay_2,az_2]//$加速度向量
$\}\}$
实际上,在机器学习和特别是深度学习模型中,对于初步定义的特征向量的集合(字典\(\mathrm{V}\))是基于球面谐函数从数据中直接提取的。这些特征向量不是简单地随机初始化,而是通过计算每个点的球面谐函数,从而捕获与SO(3)群相关的旋转等变信息。这样生成的特征向量更能精确地反映分子或几何形体在三维空间中的物理和化学性质。
2、点卷积
点卷积(Point Convolution)是一种专门为处理几何数据设计的卷积操作。在传统的卷积神经网络(CNN)中,卷积是在规则的网格上执行的(例如图像的像素网格),而点卷积则需要在不规则的点集上执行,这些点集没有固定的拓扑结构。具体来说,在标准的卷积网络中,权重共享意味着卷积核在整个输入特征图上滑动,并在每个局部区域应用相同的权重。在点卷积中,这一概念需要被重新解释,因为没有固定的局部区域可以简单地“滑动”。所以,点卷积通过定义每个点的局部邻域来操作。局部邻域可以通过找到每个点在空间中最近的\(\mathrm{k}\)个邻点来确定,或者通过在一个固定半径内搜索邻点。对于每个点,其局部邻域中的每个节点都有一个特征向量,例如第1节所讲的例子。点卷积操作涉及到将每个邻点的特征向量与一个权重矩阵相乘,这个权重矩阵对所有点是共享的,以便提取局部特征。最后,将邻居点的变换后特征聚合(例如,通过加和或最大池化)来生成一个新的特征描述符,代表中心点的局部结构。
与传统的卷积不同,点卷积除了能够处理无序和不规则的数据结构以外,还需要在处理过程中保证几何图数据的在三维空间下的对称性。这里的关键是使用Wigner-D矩阵,Wigner-D矩阵提供了一个旋转群SO(3)的不可约表示,这可以用来构建旋转等变的特征。数学上,如果\(Y^{(l)}\)是一个在SO(3)下的等变函数,则对于所有\(g\in SO(3)\)和输入向量\(\mathrm{r}\),有:
\(Y^{(l)}(R(g)\boldsymbol{r})=\sum_{m^{\prime}}D_{mm^{\prime}}^{(l)}(g)\boldsymbol{Y}_{m^{\prime}}^{(l)}(\boldsymbol{r})\)
其中\(D^{(l)}(g)\)是Wigner-D矩阵,\(\mathrm{m}\)和\(\mathrm{m~'}\)是索引这些特征向量的量子数。为了设计旋转等变的点卷积,基于上式定义的\(Y^{(l)}\),可以定义卷积核的形式为:
\(F_{crm}^{\left(l_f,l_i\right)}(r)=R_c^{l_f,l_i}(r)Y_m^{\left(l_f\right)}(\hat{r})\)
其中:
(1)\(Y_m^{(l_f)}(\hat{r})\)是在单位向量\(\hat{r}\)上计算的球面谐函数的\(\mathrm{m}\)阶分量。
(2)\(\hat{r}\)是单位长度的\(\mathbf{r}\)向量,而\(\mathrm{r}\)是\(\mathbf{r}\)的长度。
(3)\(R_c^{l_f,l_i}(r)\)是一个关于\(\mathrm{r}\)(也就是点与卷积中心之间的距离)的径向函数。
(4)\(l_{i}\)表示输入特征的旋转阶数,\(l_{f}\)表示点卷积核的旋转阶数。
(5)下标\(\mathrm{c}\)通常表示特征或点卷积核的通道。
因此,\(R_c^{l_f,l_i}(r)\)径向函数负责对每个半径\(\mathrm{r}\)的点提供一个权重,这个权重会影响卷积核如何处理距离中心点不同远近的其他点的信号。在实践中,\(R_c^{l_f,l_i}(r)\)可以是某种固定的距离函数,如cos函数、-log函数等;也可以是可学习的参数化函数,如高斯径向基函数,参数(如宽度、中心等)在训练过程中被优化。另一种常见的做法是将\(r,l_i,l_f\)和\(\mathrm{c}\)作为特征输入到一个小的神经网络(例如多层感知机MLP),该网络学习在给定这些输入的情况下输出正确的距离权重。
在进行点卷积时,传统的卷积操作不能直接应用,因为几何空间的点缺乏像图像像素那样的规则结构。因此,需要一种方法来定义点云中点的“局部区域”以及如何在这些局部区域内聚合信息。这就是\(R_c^{l_f,l_i}(r)\)和球面谐函数\(Y_{m}^{(l_{f})}(\hat{r})\)发挥作用的地方,具体步骤如下:
(1)首先选择一个中心点。
(2)对于中心点的每个邻居,计算从中心点到邻居点的向量\(\mathbf{r}\)。
(3)将这个向量规范化到单位向量\(\hat{r}\),并计算它的长度\(\mathrm{r}\)。
(4)应用球面谐函数\(Y_m^{(l_f)}(\hat{r})\)来捕捉中心点和它邻居之间的角度关系。
(5)使用\(R_c^{l_f,l_i}(r)\)函数编码距离\(\mathrm{r}\)的信息,这样近的和远的点可以有不同的影响力。
(6)定义卷积核\(F_{crm}^{\left(l_f,l_i\right)}(r)=R_c^{l_f,l_i}(r)\boldsymbol{Y}_m^{\left(l_f\right)}(\hat{\boldsymbol{r}})\),这个卷积核综合考虑的距离信息和角度关系,使用它聚合邻居的信息,形成该中心点的新特征向量。
实际上,点卷积就是在SchNet模型中提出的连续性滤波卷积的滤波器生成函数基础上,乘了一个被Winger D矩阵和球谐函数计算得到的角度分布,其他的没有区别。这样的设计额外考虑了角度分布,允许滤波器捕捉局部邻域内的方向依赖特性,从而达到旋转等变性。
3、Clebsch-Gordan coefficients
此外,在上述点卷积计算流程的最后一步(第六步)定义完卷积核后,节点\(\mathrm{i}\)的特征与其邻居节点\(\mathrm{j}\)的特征进行卷积操作更新得到新的点\(\mathrm{i}\)特征时,需要使用CG系数。因为节点\(\mathrm{i}\)的特征与其邻居节点\(\mathrm{j}\)的特征是在多种旋转阶\(\mathrm{l}\)下定义的,两个具有不同角量子数\(\mathrm{l}\)和\(\mathrm{m}\)的特征想要进行信息交互更新得到新特征时,不能简单地将它们进行相加或相乘,否则会使得新特征不能保留系统的旋转等变性。此时需要一种方法来确保合并后的新特征仍然符合旋转等变性,CG系数就是用来做这件事的。CG系数在量子力学中用于合并两个带有角动量的状态,生成一个新的状态,这个状态具有明确定义的总角动量,这种角动量守恒是实现旋转等变性的基础。在张量场神经网络中,类似的原理被应用于特征融合,使用CG系数意味着可以生成具有良好旋转特性的复合特征,这些特征可以用于训练模型,使其能够预测旋转不变的分子性质。结合CG系数后,在一个更全面的点卷积公式表述如下:
\(\mathcal{L}_{icm_{out}}^{(l_{out})}\left(\boldsymbol{r}_i,\boldsymbol{V}_{icm_{in}}^{(l_m)}\right)=\sum_{m_f,m_{in}}C_{l_{out},m_{out}}^{(l_{in},m_{in}),(l_f,m_f)}\sum_{j\in S}F_{cm_f}^{(l_f,l_{in})}\left(\boldsymbol{r}_{ij}\right)\boldsymbol{V}_{jcm_{in}}^{(l_{in})}\)
其中:
(1)\(\mathcal{L}_{icm_{out}}^{(l_{out})}\left(r_i,V_{icm_{in}}^{(l_{in})}\right)\)是经过点卷积操作后得到的新特征向量,对应于输入节点\(\mathrm{i}\)的一个特定旋转阶数\(l_{out}\)。
(2)\(r_{ij}=r_i-r_j\)表示从点\(\mathrm{i}\)到点\(\mathrm{j}\)的位置向量。
(3)\(\sum_{j\in S}F_{cm_f}^{\left(l_f,l_{in}\right)}\left(r_{ij}\right)\)代表了连续卷积核生成函数,它依赖于距离向量\(r_{ij}\)。为每个通道\(\mathrm{c}\)根据距离生成卷积核的相关权重。
(4)\(V_{jcm_{in}}^{(l_{in})}\)是点\(\mathrm{j}\)的旋转阶数\(l_{in}\)的输入特征,该特征是球谐函数编码的。
(5)\(C_{l_{out},m_{out}}^{(l_{in},m_{in}),(l_{f},m_{f})}\)是Clebsch-Gordan系数,用于处理输入特征与卷积核特征的旋转阶数合并,这个合并过程需要遵守物理规则以确保结果特征的旋转等变性,三个限制如下所示。
1.总磁量子数守恒:\(M=m_1+m_2\)
在量子系统中,总磁量子数守恒意味着在角动量合并过程中,子系统的磁量子数之和必须等于合并后系统的总磁量子数。在神经网络的背景下,这可以被看作是特征维度或特征信息在处理过程中的一种“信息守恒”。当网络层通过CG系数合并输入特征时,系数确保信息的总和在旋转变换后仍然得到保持,这使得网络的输出对于输入的旋转保持不变或适应。
2.三角不等式:\(\left|l_1-l_2\right|\leq L\leq l_1+l_2\)
三角不等式在物理上保证了角动量的可能值必须满足这一范围。在神经网络中,通过 CG系数应用这一原则,帮助模型在融合不同层或不同部分的特征时,维持一个有效的特征“规模”或“层级”,确保输出特征不会超出合适的处理范围。这对于维持模型对旋转的适应性至关重要,因为它保证了特征合并的正确性和适应性。
3.宇称守恒: \(l_1+l_2+L\)通常为偶数
在物理中,特别是在量子力学的角动量理论中,宇称守恒与宇称(parity)的概念密切相关。宇称描述了物理系统在空间坐标反演(如点\((x,y,z)\)变成\((-x,-y,-z)\)时的行为。对于轨道角动量,宇称可以通过轨道角动量的量子数\(\mathrm{l}\)来描述,其宇称为\(\left(-1\right)^l\)。这意味着,如果轨道角动量的量子数\(\mathrm{l}\)是偶数,宇称为正(不变),如果\(\mathrm{l}\)是奇数,则宇称为负(反转)。
在组合两个角动量系统时(比如两个轨道角动量,量子数分别为\(l_{1}\)和\(l_{2}\),结果的总角动量量子数为\(\text{L}\)),宇称的乘积为\((-1)^{l_1}\times(-1)^{l_2}\times(-1)^L=(-1)^{l_1+l_2+L}\)。因此,当\(l_{1}+l_{2}+L\)是偶数时,这个乘积是1(即总系统宇称不变),而当\(l_{1}+l_{2}+L\)是奇数时,宇称为-1 (总系统宇称反转)。
如果系统的物理属性在空间坐标反演后保持不变,即系统的宇称是正的,则系统被称为宇称不变的。例如,如果一个系统或状态在坐标反转后看起来与反转前相同,则这个系统或状态的宇称是不变的。另一方面,对称性通常指的是在进行某些变换(如平移、旋转、反射)后系统的物理法则保持不变的性质。因此,宇称是对称操作的一种形式,在神经网络中使用CG系数时,这种守恒原则有助于保证在旋转对称性的情况下,网络能正确处理对称与非对称特征的变换。
上述这些守恒原则为神经网络提供了一种强大的数学工具,用以保证在处理具有旋转变换的输入数据时,网络能够适当地维持输入特征的旋转特性。如果不满足上述三个守恒规则,则CG系数的值为0。
实际上,上述公式可以看作一个双线性层的计算。双线性层(Bilinear Layer)是一种通过两个输入向量的张量积生成输出的层。其一般形式为:\(z=x^{\mathrm{T}}Wy\),其中\(\mathbf{x}\)和\(\mathbf{y}\)是输入向量,\(\mathrm{W}\)是权重矩阵,可以将公式分解成几个部分:
(1)输入特征\(V_{jcm_\mathrm{in}}^{(l_\mathrm{in})}\):表示输入的特征映射,带有不可约表示\(l_{\mathrm{in}}\)和通道索引\(C_{in}\)。
(2)卷积核\(F_{cm_f}^{\left(l_f,l_{\mathrm{in}}\right)}(r_{ij})\):表示基于相对位置向量\(r_{ij}\)的卷积核。
(3)Clebsch-Gordan系数\(C_{(l_{\mathrm{out}},m_{\mathrm{out}})}^{(l_{\mathrm{in}},m_{\mathrm{in}}),(l_{f},m_{f})}\):用于将输入特征和卷积核的角动量态结合起来。
从双线性层的角度来看,输入特征\(\mathrm{V}\)和卷积核\(\text{F}\)可以分别看作双线性层的两个输入,而Clebsch-Gordan系数则类似于权重张量\(\boldsymbol{W}\):
(1)输入向量:\(x=V_{jc_{\mathrm{in}}}^{(l_{\mathrm{in}})}\)
(2)输入向量:\(\boldsymbol{y}=F_{c_{mf}}^{(l_{f},l_{\mathrm{in}})}(\boldsymbol{r}_{ij})\)
(3)权重张量:\(\boldsymbol{W}=C_{(l_{\mathrm{out}},m_{\mathrm{out}})}^{(l_{\mathrm{in}},m_{\mathrm{in}})(l_{f},m_{f})}\)
双线性层的计算方式可以看作是具有一个隐藏层的多层感知器,通过两个输入向量和一个权重张量进行加权求和,生成输出。Clebsch-Gordan 系数的作用是确保了张量积的正确性和有效性,确定哪些输入特征组合是物理上允许的,以及这些组合如何生成符合输出表示的特征,即根据物理规则决定隐藏层的权重矩阵哪些应该置0。
CG系数的计算方式依赖于固定且复杂的数学公式,其的值取决于两个原始角动量的量子数\(l_{1},m_{1}\)和\(l_{2},m_{2}\),以及结果角动量的量子数\(L,M\),如下所示:
\(\begin{aligned}&C\left(l_1,m_1;l_2,m_2;L,M\right)=\\&\delta_{m_1+m_2,M}\sqrt[]{\frac{(2L+1)(L+l_1-l_2)!(L-l_1+l_2)!(l_1+l_2-L)!}{\left(l_1+l_2+L+1\right)!}}\sqrt{\frac{(L+M)!(L-M)!}{\left(l_1+m_1\right)!\left(l_1-m_1\right)!\left(l_2+m_2\right)!\left(l_2-m_2\right)!}}\times f\end{aligned}\)
其中,\(\mathrm{f}\)是一个涉及Wigner 3j符号的表达式,其具体形式为:
\(f=(-1)^{l_1-l_2+M}\begin{pmatrix}l_1&l_2&L\\m_1&m_2&-M\end{pmatrix}\)
\(\left.\left(\begin{array}{ccc}{l_{1}}&{l_{2}}&{L}\\{m_{1}}&{m_{2}}&{-M}\end{array}\right.\right)\)是Wigner 3j符号,它反映了角动量合成的几何和对称性属性,并直接涉及到角动量的选择规则和合成。在实际应用中,这些系数通常使用数学软件如 Mathematica、Python 中的 SymPy 库等来计算,以表格的形式存储,以便快速查找和使用,无需每次使用时重新计算,如表1所示。
在神经网络或其他计算应用中,通常会预先加载一个包含所有所需\(l_{1},l_{2},l\)组合的CG系数的数据库或文件。这样做可以避免计算上的延迟,并提高效率。在软件中,这些系数可能已经以优化的数据结构存储,以便快速访问和使用。
4、自相互作用层
TFN通过上一小节介绍的点卷积层来聚合某节点的邻居信息,更新得到新的节点表示。此外,TFN还通过子相互作用层来仅更新某节点自身的特征向量(不考虑邻居信息),自相互作用层可以看作类似于计算机视觉中的1×1的2D卷积操作,其公式表示如下:
\(V_{\mathrm{acm}}^{(l+1)}=\sum_{c^{\prime}}W_{\mathrm{cc}}^{(l)}V_{\mathrm{ac}^{\prime}\mathrm{m}}^{(l)}\)
该公式可以解释为:\(V_{\mathrm{ac^{\prime}m}}^{(l)}\)是层\(\mathrm{l}\)中节点\(\mathrm{a}\)经球谐函数编码的输入特征向量,其中\(\mathrm{m}\)表示特征向量在张量场中的一个特定分量,\(c^{\prime}\)表示输入通道。\(W_{cc^{\prime}}^{(c)}\)是权重矩阵,它定义了从输入通道\(c^{\prime}\)到输出通道\(\mathrm{c}\)的线性变换方式。总和\(\sum_{c^{\prime}}\)表示对所有输入通道的累加,这允许模型从不同的输入通道学习到不同的特征,并将它们合并成新的节点表示。该自相互作用层的输出可以表示为\(V_{\mathrm{acm}}^{(l+1)}\),即节点\(\mathrm{a}\)在层\(\text{l+1}\)的特征向量,其每个通道\(\mathrm{c}\)的值由上述线性组合计算得到。这样的操作使得每个节点能够独立更新其特征表示,而不受邻居节点的直接影响,这对于保持特征的局部性非常重要。在进一步的网络层中,这些更新过的特征向量\(V_{\mathrm{acm}}^{(l+1)}\)可以被用作输入到后续的SO(3)卷积层或其他类型的层,以实现更复杂的特征提取和变换。这种设计使TFN能够适应各种各样的数据结构和任务,提供灵活而强大的模型架构以处理具有复杂几何和拓扑结构的数据。
5、激活函数
与传统神经网络一样,TFNs中的卷积层后通常会跟随激活函数。这些非线性函数是必需的,因为它们允许网络学习和表示复杂的非线性关系。在处理分子数据时,合适的激活函数可以帮助模型识别分子内部的化学反应路径和能量变化等复杂现象。需要注意的是,在处理旋转阶数为0和大于0两种情况时,激活函数的定义是不同的。具体来说对于旋转阶数为0的特征,通常可以应用标准的非线性激活函数,如ReLU、Sigmoid或Tanh,这些激活函数能够处理标量输入,适用于不受旋转变换影响的特征。
然而,对于旋转阶数大于0的情况,即那些在旋转下变换的特征,标准的激活函数不能直接应用,因为这些特征具有方向性,需要保持其在空间旋转下的一致性。为了保持旋转等变性,激活函数在这种情况下只应用于特征向量的模长,因为模长是旋转不变的。在旋转阶数大于0的情况下,用非线性函数处理后的模长来缩放原始特征向量的每个分量。这个过程可以想象成对原始向量的长度进行缩放,但是保持其指向不变。通过这样的操作,特征向量在经过非线性变换后,仍然保持对旋转的敏感性。这样,当整个系统旋转时,变换后的特征向量仍然以一种一致的方式表达同一个物理现象或结构特征。
6、顶层结构
在TFNs中,顶层设计包含池化层和全连接层。池化层的职能是整合整个分子的信息。通过对所有原子的高维表示进行汇总,网络可以获得一个全局描述,这对于预测分子级别的性质(如分子的总能量或偶极矩)是必要的。这个过程可能包括简单的操作,比如计算所有原子特征的平均值,或者更复杂的加权和,这些加权可能依赖于原子之间的相对位置。
在经过一系列卷积和池化操作之后,TFN可以包括一或多个全连接层。这些层可以进一步组合和转换全局描述,以产生最终的输出,如分子的性质预测。全连接层能够学习不同原子特征之间复杂的相互作用,并将其映射到所需的输出空间。
最终,TFN的输出层将网络的高维表示映射到所关心的预测上。在分子动力学模拟中,这可能是能量级别、力的大小和方向、电荷分布等。这个层通常会经过训练来最小化预测和实际值之间的差距,使用的是回归分析中常见的损失函数,如均方误差。
整个TFN架构通过监督学习来训练。给定一组分子结构和相应的目标属性(如实验测量的能量),网络通过优化算法(如梯度下降)调整其权重,以最小化预测输出和目标输出之间的差异。在这个过程中,网络学习到了如何精确捕捉分子内部的复杂相互作用,以及这些相互作用如何影响分子的宏观性质。
7、TFN的消息传递机制
TFN的消息传递过程可以分为以下三个步骤:
1、消息传递与更新:
第一步是节点之间的信息传递和更新。具体的数学表示如下:
\(\boldsymbol{M}_{ij}^{(L)}=Y^{(L)}\left(\frac{\boldsymbol{x}_{ij}}{\|\boldsymbol{x}_{ij}\|}\right)\odot\boldsymbol{W}_{\mathrm{CG}}\boldsymbol{V}_j^{(L)}\)
其中:
(1) \(x_{ij}=x_i-x_j\)表示节点\(\mathrm{i}\)和节点\(\mathrm{j}\)之间的相对位置向量;
(2)\(\boldsymbol{W}\)是一个可学习的权重矩阵,并遵守于CG系数的限制;
(3)\(Y(|x_{ij}|)\)是一个依赖于相对位置向量\(\mathcal{X}_{ij}\)的球谐基函数;
(4)\(\odot\)表示元素逐位相乘;
(5)\(V_j^{(L)}\)是节点\(\mathrm{j}\)在第\(\mathrm{L}\)层的特征向量。
在上述公式中,\(\boldsymbol{M}_{ij}^{(L)}\)表示从节点\(\mathrm{j}\)传递到节点\(\mathrm{i}\)的消息,这个消息通过节点\(\mathrm{j}\)的特征\(V_j^{(L)}\)和一个基于相对位置向量的球谐函数\(\text{Y}\)进行加权平均后得到。
接下来,节点\(\mathrm{i}\)在第\(\mathrm{L}\)层的更新特征向量\(U_i^{(L)}\)通过以下公式计算:
\(U_i^{(L)}=V_i^{(L)}+\sum_{j\in\mathcal{N}(i)}M_{ij}^{(L)}\)
其中:
(1) \(U_i^{(L)}\)表示节点\(\mathrm{i}\)在第\(\mathrm{L}\)层的更新特征向量;
(2) \(V_{i}^{(L)}\)表示节点\(\mathrm{i}\)在第\(\mathrm{L}\)层的初始特征向量;
(3)\(\mathcal{N}(i)\)是节点\(\mathrm{i}\)的邻居节点集合。
通过上述公式,可以将来自邻居节点的消息\(\boldsymbol{M}_{ij}^{(L)}\)聚合到当前节点的特征\(\tilde{V}_i^{(L)}\)上,实现特征的更新和信息的传播。
2、自我交互:
第二步是对聚合后的节点特征进行自我交互,公式表示如下:
\(\tilde{V}_i^{(L)}=\{U_i^{(L)}W^{(L)}\}_{\ell\in L}\)
其中,\(W^{(L)}\)是一个可学习的权重矩阵。自我交互是对特征向量进行的线性变换,在这个步骤中,每个节点的特征通过自我交互的权重矩阵\(W^{(L)}\)进行变换,从而捕捉到特征之间的线性关系。
3、非线性化(激活函数):
消息传递过程的第三步是对节点特征进行非线性化处理。这个步骤的目的是引入非线性因素,使得模型能够捕捉到更加复杂的特征和关系。具体的数学表示如下:
\(\tilde{\boldsymbol{V}}_i^{(L)}=\left\{\tilde{\boldsymbol{V}}_i^{(\ell)}\sigma\left(\left\|\tilde{\boldsymbol{V}}_i^{(\ell)}\right\|_2+b_i^{(\ell)}\right)\right\}_{\ell\in L}\)
其中:
(1)激活函数\(\sigma\):激活函数是深度学习模型中引入非线性的重要组件。常见的激活函数有ReLU和Sigmoid等。ReLU将负值映射为0,保留正值,而Sigmoid则将值映射到0和1之间。通过引入非线性激活函数,模型能够更好地拟合复杂的非线性关系。
(2)特征向量的模长:计算节点\(\text{i}\)在第\(\ell\)层特征向量的L2范数(欧几里得范数)即为该特征向量的模长。通过计算模长,可以获得特征向量的大小信息,这对于后续的缩放和归一化操作至关重要。
(3)偏置项\(b_{i}^{(\ell)}\):偏置项是一个可学习的参数,用于调整特征向量的模长,使得模型能够更灵活地适应不同的特征分布。
(4)非线性化的意义:通过非线性化处理,可以引入复杂的特征关系,提升模型的表达能力。非线性化后的结果可以视为对节点\(\begin{array}{c}{i}\end{array}\)的特征进行缩放,使其更适应模型的学习需求。具体的计算结果表示为:\(\sigma\left(\left|\tilde{V}_i^{(\ell)}\right|_2+b_i^{(\ell)}\right)\)。这一结果可以看作是一个缩放因子,对节点\(\begin{array}{c}{i}\end{array}\)的特征\(\tilde{V}_i^{(\ell)}\)进行缩放,从而使得特征更具辨识度和适应性。
通过以上步骤,TFN能够在特征传递和更新过程中有效地捕捉到节点之间的复杂关系,并通过非线性化处理增强模型的表达能力和拟合效果。这一步骤在整个消息传递过程中起到了至关重要的作用,确保了模型的灵活性和鲁棒性。
作者
arwin.yu.98@gmail.com