[论文解读] Molecular Geometry-aware Transformer for accurate 3D Atomic System modeling

Zheng Yuan, Yaoyun Zhang+4 more|arXiv (Cornell University)|Feb 2, 2023

Machine Learning in Materials ScienceMaterials Science被引用 3

一句话总结

Moleformer 提出了一种几何感知的 Transformer 架构，通过基于相对距离和角度的旋转与平移不变空间编码，对原子、化学键和非键合原子对进行建模。该方法在 OC20 数据集上实现了松弛能量预测的最先进性能，在 QM9 上也取得了具有竞争力的结果，通过显式建模高阶几何相互作用，超越了以往的 GNN 和 Transformer 模型。

ABSTRACT

Molecular dynamic simulations are important in computational physics, chemistry, material, and biology. Machine learning-based methods have shown strong abilities in predicting molecular energy and properties and are much faster than DFT calculations. Molecular energy is at least related to atoms, bonds, bond angles, torsion angles, and nonbonding atom pairs. Previous Transformer models only use atoms as inputs which lack explicit modeling of the aforementioned factors. To alleviate this limitation, we propose Moleformer, a novel Transformer architecture that takes nodes (atoms) and edges (bonds and nonbonding atom pairs) as inputs and models the interactions among them using rotational and translational invariant geometry-aware spatial encoding. Proposed spatial encoding calculates relative position information including distances and angles among nodes and edges. We benchmark Moleformer on OC20 and QM9 datasets, and our model achieves state-of-the-art on the initial state to relaxed energy prediction of OC20 and is very competitive in QM9 on predicting quantum chemical properties compared to other Transformer and Graph Neural Network (GNN) methods which proves the effectiveness of the proposed geometry-aware spatial encoding in Moleformer.

研究动机与目标

• 为解决现有 Transformer 模型在 3D 原子系统建模中的局限性，即仅将原子作为输入，未能显式建模化学键、键角和非键合相互作用。
• 通过将键角和二面角等几何因素直接整合到注意力机制中，提升分子能量和性质预测的准确性。
• 通过引入周期性边界条件校正和数据增强，提升模型在分布外分子系统中的泛化能力。
• 通过在键长上施加边级别的辅助损失，验证其在几何监督下对能量预测性能的提升效果。

提出的方法

• Moleformer 以原子和边（化学键与非键合原子对）作为输入，扩展了以往仅使用原子输入的 Transformer 模型。
• 提出一种几何感知的空间编码方法，计算节点与边之间的相对距离和角度，确保旋转与平移不变性。
• 在节点和边层面均应用自注意力机制，使每一层能够显式建模三体和四体相互作用（如键角、二面角）。
• 引入一种新颖的周期性边界条件（PBC）校正方法，提升具有周期性结构系统的建模精度，尤其适用于催化剂模拟。
• 采用噪声节点数据增强方法，提升在 OC20 数据集中分布外吸附物上的泛化能力。
• 引入边级别的辅助损失，对预测的键长进行监督，从而提升能量估计性能。

Figure 1: The overview of Moleformer. (a) A toy example of a catalyst system where \ch CH3 is the adsorbate and the surface is composed of copper atoms. Moleformer prefers edges between adosrbate-adsorbate atoms and adsorbate-surface atoms as model inputs. (b) Architecture of Moleformer: Nodes and s

实验结果

研究问题

• RQ1 在 Transformer 架构中显式建模化学键和非键合原子对，是否能相比仅基于原子的模型，提升 3D 原子系统能量预测性能？
• RQ2 基于相对距离和角度的几何感知空间编码，对分子能量和性质预测任务的性能有何影响？
• RQ3 在具有周期性结构的催化系统中，引入周期性边界条件校正能在多大程度上提升能量预测精度？
• RQ4 在键长预测上增加边级别的辅助损失，是否能增强模型准确估计系统能量的能力？
• RQ5 Moleformer 在分布外分子上的泛化能力如何，特别是在 OC20 数据集中稀有吸附物的情境下？

主要发现

• Moleformer 在 OC20 数据集上实现了初始态到松弛态能量预测的最先进性能，于分布内划分的测试集中能量 MAE 达到 413 meV。
• 引入周期性边界条件校正后，IS2RS 任务的位置 MAE 降低，IS2RE 划分中能量 MAE 减少 0.014 eV，EWT 提升 0.7%。
• 若不使用噪声节点数据增强，分布外吸附物划分中的能量 MAE 显著上升，表明数据增强有效提升了泛化能力。
• 移除边级别的辅助损失后，分布内与分布外划分中的能量 MAE 均上升，证明其在几何监督中的有效性。
• 在 QM9 数据集上，Moleformer 在预测量子化学性质方面表现具有竞争力，包括分子轨道能级，优于近期的 SOTA 方法如 SEGNN 和 EQGAT。
• 消融实验证实，所有组件——PBC 校正、数据增强和边级别损失——均对性能有贡献，其中边级别损失带来的提升最大。

Figure 2: PBC Correction. (a)-(c) An example of a catalyst system where \ch CH3 is the adsorbate and the surface is composed of copper atoms viewed along the <110>orientation. The dotted lines show the borders of the unit cell. Grey-colored small atom is a hydrogen atom that goes across the unit cel