基于机器学习的方法在预测分子能量和性质方面表现出很强的能力。分子能量至少与原子、键、键角、扭转角和非键原子对有关。以前的Transformer模型只使用原子作为输入，缺乏对上述因素的显式建模。为了减轻这种限制，作者提出了Moleformer，这是一种新颖的Transformer架构，它将节点（原子）和边（键和非键原子对）作为输入，并使用旋转平移不变的几何感知空间编码来模拟它们之间的相互作用。空间编码计算节点和边之间的相对位置信息，包括距离和角度。作者在OC20和QM9数据集上对Moleformer进行了基准测试。

来自：Molecular Geometry-aware Transformer for accurate 3D Atomic System modeling

背景概述

分子系统能量力场的基本形式可以分解为基于键、键角、扭转角和非键相互作用等因素的函数。为了包含上述因素并进行能量估计，神经网络需要学习原子-原子（键和非键原子对）、原子-键（键角）和键-键（扭角）之间的相互作用。而过去基于Transformer的模型只学习原子-原子之间的相互作用。

为了扩展信息并利用Transformer，作者提出Moleformer，它显式模拟原子-原子对之间的相互作用。Moleformer将原子、键和非键原子对作为输入，并应用平移旋转不变几何感知空间编码来捕获这些输入之间的几何关系。空间编码是通过原子-原子对的相对距离和角度来计算的，这足以确定原子-原子对之间的相对空间几何信息。原子-键之间的相互作用包括力场中的键角，键-键之间的相互作用包括力场中的扭转角。

方法

对于一个3D原子系统（$N$个原子），由核电荷（nuclear charges） { z 1 , . . . , z N } \left\{z_{1},...,z_{N}\right\} {z1​,...,zN​}和3D坐标$\left\{x_1,...,x_N\in R^3\right\}$描述。目标是预测系统的性质$y\in R$，可以是系统能量，偶极矩和热容等。

Moleformer（图1b）是一种旋转平移不变的Transformer，用于3D分子能量和性质预测。Moleformer使用节点（原子）和选定的边（包括成键和非成键原子对）作为输入。如果一个原子系统包含$N$个原子，则有$C_N^2$个可能的边。将所有可能的边都包含到模型中会消耗很大计算量，并且可能引入无信息的边。作者选择 Top-M 距离最近的原子对作为边。对于像催化剂体系这样由吸附质（adsorbate）和表面（surface）组成的原子体系（图1a），可以选择吸附质-吸附质原子对和吸附质-表面原子对作为边，而不是表面-表面原子对。原因是这样的原子对之间的边在边长和原子类型上更加多样化，在能量建模中可能包含更多的信息。因此需要为表面-表面原子对的距离进行惩罚，以更好选择边。

---

表面原子之间的相互作用较频繁，而吸附质原子与表面原子的相互作用较小，不平衡其重要性会学习到冗余的知识。

---

对于输入embedding，作者用核电荷$z_i,z_j$和两原子间的距离$||r_{ij}||=||x_i-x_j||$编码边$(i,j)$：$$h_{ij}^0=RBF(z_i,z_j,||r_{ij}||)$$原子$i$则编码为：$$h_i^0=Embed(z_i)+W_R\sum _j(RBF(z_i,z_j,||r_{ij}||))+b_R$$其中，$W_R$和$b_R$是线性层的参数，用于对齐$Embed$和RBF的维数。

将${\left\{h_i^0\right\}}_{i=1}^N$和$\left\{h_{ij}^0\right\}$拼接作为输入，并应用$L$层geometry-aware Transformer block建模原子和边。对于第$l$层，输入为${\left\{h_i^l\right\}}_{i=1}^N$和$\left\{h_{ij}^l\right\}$。

Transformer通常通过在输入表示中添加位置嵌入来编码绝对位置信息（即词序）。对于三维原子建模，分子的表示是旋转平移不变的，因此需要相对位置信息。作者在每个几何感知Transformer层中通过在自注意力计算中添加一个偏差来编码相对位置信息：$$A_{ij}=\frac{(h_i^l{W}_q)(h_j^l{W}_k{)}^T}{\sqrt{d}}+b(i,j)$$$$a_{ij}=\frac{exp{A}_{ij}}{{\sum }_{k}exp{A}_{ik}}$$$$h_i^{l+1}=\sum _j(h_i^l{W}_v)a_{ij}$$其中，$i,j,k$可以表示一个节点或者一个边，$b(i,j)$为作者提出的空间编码。

特别的，如果$i,j$都是原子，唯一的相对几何信息就是原子之间的距离（图1c），因此，作者使用RBF计算$b(i,j)$：$$b(i,j)=RBF(z_i,z_j,||r_{ij}||)$$边$(i,j)$和节点$k$来自三角形$i,j,k$，相对几何信息由两条边的长度$||r_{ik}||$和$||r_{jk}||$以及转角$\angle ikj$决定（图1d）。$b(ij,k)$为：$$b(ij,k)=RBF(||r_{ik}||)+RBF(||r_{jk}||)+RBF(cos\angle ikj)$$连接不同节点对的两条边$(i,j)$和$(k,l)$构成一个四边形，相对几何信息由六条边的长度决定，比如$(i,j),(i,k),(i,l)(j,k),(j,l),(k,l)$。边$(i,j)$和$(k,l)$编码为$h_{ij}^l$和$h_{kl}^l$。作者使用另外四条边编码$b(ij,kl)$。此外，通过边$(i,j)$和$(k,l)$的角度增强空间编码（图1e）：$$b(ij,kl)=RBF(||r_{ik}||)+RBF(||r_{jk}||)+RBF(||r_{il}||)+RBF(||r_{jl}||)+RBF(cos(r_{ij},r_{kl}))$$节点$h_i$和边$h_{ij}$的最终表示是通过重复将表示发送到Transformer层中获得。该方法在不增加参数数量的情况下提高了模型的性能。

• 图1a：一个催化剂系统的例子，其中CH3是吸附质，表面由铜原子组成。Moleformer倾向于将吸附-吸附原子和吸附-表面原子之间的边作为模型输入。
• 图1b：Moleformer架构：节点和所选边作为输入。
• 图1c-e：分别展示了几何感知空间编码如何在Transformer中的节点-节点、边缘-节点和边缘-边缘之间相互作用。虚线表示边缘，黄色实线表示空间编码中使用的几何信息。

受到包括键作为分子系统能量因子的启发，不仅使用节点表示，还使用边缘表示进行能量（或其他标量属性）预测：$$y=\sum _{i}MLP(h_i)+\sum _{ij}MLP(h_{ij})$$其中MLP是一个非线性的2层感知机。