Equiformer: Equivariant Graph Attention Transformer for 3D Atomistic Graphs

Paper | OpenReview | Poster | Slides

This repository contains the official PyTorch implementation of the work "Equiformer: Equivariant Graph Attention Transformer for 3D Atomistic Graphs" (ICLR 2023 Spotlight).

Additionally, in our subsequent work, we find that we can generalize self-supervised learning similar to BERT, which we call DeNS (Denoising Non-Equilibrium Structures), to 3D atomistic systems to improve the performance of Equiformer on MD17 dataset. We provide the implementation of training Equiformer with DeNS on MD17 below. Please refer to the paper and the code for further details.

Content

0. Environment Setup
1. Training
2. File Structure
3. Acknowledgement
4. Citation

Environment Setup

Environment

See here for setting up the environment.

QM9

The dataset of QM9 will be automatically downloaded when running training.

MD17

The dataset of MD17 will be automatically downloaded when running training.

OC20

The dataset for different tasks can be downloaded by following instructions in their GitHub repository. After downloading, place the datasets under datasets/oc20/ by using ln -s. Take is2re as an example:

    cd datasets
    mkdir oc20
    cd oc20
    ln -s ~/ocp/data/is2re is2re

Training

QM9

1. We provide training scripts under scripts/train/qm9/equiformer. For example, we can train Equiformer for the task of alpha by running:

       sh scripts/train/qm9/equiformer/target@1.sh

2. The QM9 dataset will be downloaded automatically as we run training for the first time.

3. The target number for different regression tasks can be found here.

4. We also provide the code for training Equiformer with linear messages and dot product attention. To train Equiformer with linear messages, replace --model-name 'graph_attention_transformer_nonlinear_l2' with --model-name 'graph_attention_transformer_l2' in training scripts.

5. The training scripts for Equiformer with linear messages and dot product attention can be found in scripts/train/qm9/dp_equiformer.

6. Training logs of Equiformer can be found here.

MD17

1. We provide training scripts under scripts/train/md17/equiformer. For example, we can train Equiformer for the molecule of aspirin by running:

       sh scripts/train/md17/equiformer/se_l2/target@aspirin.sh    # L_max = 2
       sh scripts/train/md17/equiformer/se_l3/target@aspirin.sh    # L_max = 3

2. Training logs of Equiformer with $L_{max} = 2$ and $L_{max} = 3$ can be found here ($L_{max} = 2$) and here ($L_{max} = 3$). Note that the units of energy and force are kcal mol $^{-1}$ and kcal mol $^{-1}$ Å $^{-1}$ and that we report energy and force in units of meV and meV Å $^{-1}$ in the paper.

3. We provide the scripts of training Equiformer with DeNS (Denoising Non-Equilibrium Structures) under scripts/train/md17/equiformer_dens. For example, we train Equiformer with DeNS for the molecule of aspirin by running:

       sh scripts/train/md17/equiformer_dens/se_l2/target@aspirin.sh    # L_max = 2
       sh scripts/train/md17/equiformer_dens/se_l3/target@aspirin.sh    # L_max = 3

4. The logs of training Equiformer with $L_{max} = 2$ and $L_{max} = 3$ with DeNS can be found here ($L_{max} = 2$) and here ($L_{max} = 3$). Note that the units of energy and force are kcal mol $^{-1}$ and kcal mol $^{-1}$ Å $^{-1}$ and that we report energy and force in units of meV and meV Å $^{-1}$ in the paper.

OC20

1. We train Equiformer on IS2RE data only by running:

       sh scripts/train/oc20/is2re/graph_attention_transformer/l1_256_nonlinear_split@all_g@2.sh

   a. This requires 2 GPUs and results in energy MAE of around 0.5088 eV for the ID sub-split of the validation set.

   b. Pretrained weights and training logs can be found here.

2. We train Equiformer on IS2RE data with IS2RS auxiliary task and Noisy Nodes data augmentation by running:

       sh scripts/train/oc20/is2re/graph_attention_transformer/l1_256_blocks@18_nonlinear_aux_split@all_g@4.sh

   a. This requires 4 GPUs and results in energy MAE of around 0.4156 eV for the ID sub-split of the validation set.

   b. Pretrained weights and training logs can be found here.

File Structure

We have different files and models for QM9, MD17 and OC20.

General

1. nets includes code of different network architectures for QM9, MD17 and OC20.
2. scripts includes scripts for training models for QM9, MD17 and OC20.

QM9

1. main_qm9.py is the training code for QM9 dataset.

MD17

1. main_md17.py is the code for training and evaluation on MD17 dataset.
2. main_md17_dens.py extends main_md17.py so that we can train Equiformer with DeNS.

OC20

Some differences are made to support:

• Removing weight decay for certain parameters specified by no_weight_decay. One example is here.
• Cosine learning rate.

1. main_oc20.py is the code for training and evaluating.
2. oc20/trainer contains the code for energy trainers.
3. oc20/configs contains the config files for IS2RE.

Acknowledgement

Our implementation is based on PyTorch, PyG, e3nn, timm, ocp, SEGNN and TorchMD-NET.

Citation

If you use our code or method in your work, please consider citing the following:

@inproceedings{
    liao2023equiformer,
    title={Equiformer: Equivariant Graph Attention Transformer for 3D Atomistic Graphs},
    author={Yi-Lun Liao and Tess Smidt},
    booktitle={International Conference on Learning Representations},
    year={2023},
    url={https://openreview.net/forum?id=KwmPfARgOTD}
}

If DeNS is helpful to your work, please consider citing the following as well:

@article{
    DeNS,
    title={Generalizing Denoising to Non-Equilibrium Structures Improves Equivariant Force Fields}, 
    author={Yi-Lun Liao and Tess Smidt and Muhammed Shuaibi* and Abhishek Das*},
    journal={arXiv preprint arXiv:2403.09549},
    year={2024}
}

Please direct any questions to Yi-Lun Liao (ylliao@mit.edu).