[Paper Review] 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 Science3 citations

TL;DR

Moleformer proposes a geometry-aware Transformer architecture that models atoms, bonds, and nonbonding atom pairs using rotational and translational invariant spatial encoding based on relative distances and angles. It achieves state-of-the-art performance on OC20 for relaxed energy prediction and competitive results on QM9, outperforming prior GNN and Transformer models through explicit high-order geometric interaction modeling.

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.

Motivation & Objective

• To address the limitation of existing Transformer models in 3D atomic system modeling, which treat only atoms as inputs and fail to explicitly model bonds, angles, and nonbonding interactions.
• To improve molecular energy and property prediction accuracy by incorporating geometric factors such as bond angles and torsion angles directly into the attention mechanism.
• To enhance generalization in out-of-distribution molecular systems by introducing periodic boundary condition correction and data augmentation.
• To demonstrate the effectiveness of edge-level auxiliary loss for improving energy prediction through geometric supervision on bond lengths.

Proposed method

• Moleformer uses atoms and edges (bonds and nonbonding atom pairs) as inputs, extending beyond standard atom-only inputs in prior Transformers.
• It introduces a geometry-aware spatial encoding that computes relative distances and angles between nodes and edges, ensuring rotational and translational invariance.
• The model applies self-attention over both nodes and edges, enabling explicit modeling of three- and four-body interactions (e.g., bond angles, torsion angles) in each layer.
• A novel periodic boundary condition (PBC) correction is applied to improve accuracy in systems with periodic structures, particularly in catalyst simulations.
• Noisy node data augmentation is used to improve generalization on out-of-distribution adsorbates in the OC20 dataset.
• An auxiliary edge-level loss is introduced to supervise predicted bond lengths, enhancing energy estimation performance.

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

Experimental results

Research questions

• RQ1 Can explicit modeling of bonds and nonbonding atom pairs in a Transformer architecture improve 3D atomic system energy prediction compared to atom-only models?
• RQ2 How does geometry-aware spatial encoding based on relative distances and angles affect performance on molecular energy and property prediction tasks?
• RQ3 To what extent does incorporating periodic boundary condition correction improve energy prediction in catalytic systems with periodic structures?
• RQ4 Does adding an edge-level auxiliary loss for bond length prediction enhance the model’s ability to estimate system energy accurately?
• RQ5 How does Moleformer generalize to out-of-distribution molecules, especially in the context of rare adsorbates in the OC20 dataset?

Key findings

• Moleformer achieves state-of-the-art performance on the OC20 dataset for initial state to relaxed energy prediction, with an energy MAE of 413 meV on the in-domain split.
• The inclusion of periodic boundary condition correction reduces the position MAE in IS2RS and improves energy MAE by 0.014 eV and EWT by 0.7% on the IS2RE split.
• Without noisy node data augmentation, the energy MAE on the out-of-domain adsorbate split increases significantly, indicating improved generalization through data augmentation.
• Removing the edge-level auxiliary loss increases energy MAE in both in-domain and out-of-domain splits, demonstrating its effectiveness in geometric supervision.
• On the QM9 dataset, Moleformer achieves competitive results in predicting quantum chemical properties, including molecular orbital energy, outperforming recent SOTA methods like SEGNN and EQGAT.
• Ablation studies confirm that all components—PBC correction, data augmentation, and edge-level loss—contribute to performance, with the edge-level loss yielding the largest improvement.

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