# Equiformer: Equivariant Graph Attention Transformer for 3D Atomistic Graphs


We find that we can generalize self-supervised learning similar to BERT, which we call **DeNS** (**De**noising **N**on-Equilibrium **S**tructures), 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.



<p align="center">
	<img src="fig/equiformer.png" alt="photo not available" width="80%" height="80%">
</p>



## Content ##
0. [Environment Setup](#environment-setup)
0. [Training](#training)
0. [File Structure](#file-structure)



## Environment Setup ##


### Environment 

See [here](docs/env_setup.md) for setting up the environment.

### MD17

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



## Training ##


### MD17

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

    ```bash
        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. We provide the scripts of training Equiformer with **DeNS** (**De**noising **N**on-Equilibrium **S**tructures) under [`scripts/train/md17/equiformer_dens`](scripts/train/md17/equiformer_dens).
For example, we train Equiformer with DeNS for the molecule of `aspirin` by running:
    ```bash
        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
    ```    



## File Structure ##

### General 

1. [`nets`](nets) includes code of different network architectures for QM9, MD17 and OC20.
2. [`scripts`](scripts) includes scripts for training models for QM9, MD17 and OC20.

### MD17

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