{"id":13699099,"url":"https://github.com/atomicarchitects/equiformer","last_synced_at":"2026-01-17T12:25:15.975Z","repository":{"id":92141933,"uuid":"607428034","full_name":"atomicarchitects/equiformer","owner":"atomicarchitects","description":"[ICLR 2023 Spotlight] Equiformer: Equivariant Graph Attention Transformer for 3D Atomistic Graphs","archived":false,"fork":false,"pushed_at":"2025-02-11T19:52:36.000Z","size":4142,"stargazers_count":251,"open_issues_count":11,"forks_count":48,"subscribers_count":6,"default_branch":"master","last_synced_at":"2025-09-09T12:44:40.177Z","etag":null,"topics":["ai-for-science","catalyst-design","computational-chemistry","deep-learning","drug-discovery","e3nn","equivariant-graph-neural-network","equivariant-neural-networks","force-fields","geometric-deep-learning","graph-neural-networks","interatomic-potentials","machine-learning","materials-science","molecular-dynamics","pytorch"],"latest_commit_sha":null,"homepage":"https://arxiv.org/abs/2206.11990","language":"Python","has_issues":true,"has_wiki":null,"has_pages":null,"mirror_url":null,"source_name":null,"license":"mit","status":null,"scm":"git","pull_requests_enabled":true,"icon_url":"https://github.com/atomicarchitects.png","metadata":{"files":{"readme":"README.md","changelog":null,"contributing":null,"funding":null,"license":"LICENSE","code_of_conduct":null,"threat_model":null,"audit":null,"citation":null,"codeowners":null,"security":null,"support":null,"governance":null,"roadmap":null,"authors":null,"dei":null,"publiccode":null,"codemeta":null,"zenodo":null}},"created_at":"2023-02-28T00:21:30.000Z","updated_at":"2025-09-08T08:34:36.000Z","dependencies_parsed_at":"2024-07-18T12:30:52.747Z","dependency_job_id":"40aa5f62-6440-4e87-a989-a9752926fc53","html_url":"https://github.com/atomicarchitects/equiformer","commit_stats":null,"previous_names":[],"tags_count":0,"template":false,"template_full_name":null,"purl":"pkg:github/atomicarchitects/equiformer","repository_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/atomicarchitects%2Fequiformer","tags_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/atomicarchitects%2Fequiformer/tags","releases_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/atomicarchitects%2Fequiformer/releases","manifests_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/atomicarchitects%2Fequiformer/manifests","owner_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners/atomicarchitects","download_url":"https://codeload.github.com/atomicarchitects/equiformer/tar.gz/refs/heads/master","sbom_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/atomicarchitects%2Fequiformer/sbom","scorecard":null,"host":{"name":"GitHub","url":"https://github.com","kind":"github","repositories_count":286080680,"owners_count":28508464,"icon_url":"https://github.com/github.png","version":null,"created_at":"2022-05-30T11:31:42.601Z","updated_at":"2026-01-17T11:50:55.898Z","status":"ssl_error","status_checked_at":"2026-01-17T11:50:55.569Z","response_time":85,"last_error":"SSL_connect returned=1 errno=0 peeraddr=140.82.121.6:443 state=error: unexpected eof while reading","robots_txt_status":"success","robots_txt_updated_at":"2025-07-24T06:49:26.215Z","robots_txt_url":"https://github.com/robots.txt","online":false,"can_crawl_api":true,"host_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub","repositories_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories","repository_names_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repository_names","owners_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners"}},"keywords":["ai-for-science","catalyst-design","computational-chemistry","deep-learning","drug-discovery","e3nn","equivariant-graph-neural-network","equivariant-neural-networks","force-fields","geometric-deep-learning","graph-neural-networks","interatomic-potentials","machine-learning","materials-science","molecular-dynamics","pytorch"],"created_at":"2024-08-02T19:00:58.176Z","updated_at":"2026-01-17T12:25:15.793Z","avatar_url":"https://github.com/atomicarchitects.png","language":"Python","funding_links":[],"categories":["Representation Learning","🔬 Domain-Specific Applications"],"sub_categories":["🌌 Physics \u0026 Astronomy"],"readme":"# Equiformer: Equivariant Graph Attention Transformer for 3D Atomistic Graphs\n\n\n**[Paper](https://arxiv.org/abs/2206.11990)** | **[OpenReview](https://openreview.net/forum?id=KwmPfARgOTD)** | **[Poster](docs/equiformer_poster.pdf)** | **[Slides](docs/equiformer_slides.pdf)**\n\n\nThis repository contains the official PyTorch implementation of the work \"Equiformer: Equivariant Graph Attention Transformer for 3D Atomistic Graphs\" (ICLR 2023 Spotlight).\n\nAdditionally, in our subsequent [work](https://arxiv.org/abs/2403.09549), 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.\nWe provide the implementation of training Equiformer with DeNS on MD17 below.\nPlease refer to the [paper](https://arxiv.org/abs/2403.09549) and the [code](https://github.com/atomicarchitects/DeNS) for further details.\n\n\n\n\u003cp align=\"center\"\u003e\n\t\u003cimg src=\"fig/equiformer.png\" alt=\"photo not available\" width=\"80%\" height=\"80%\"\u003e\n\u003c/p\u003e\n\n\n\n## Content ##\n0. [Environment Setup](#environment-setup)\n0. [Training](#training)\n0. [File Structure](#file-structure)\n0. [Acknowledgement](#acknowledgement)\n0. [Citation](#citation)\n\n\n\n## Environment Setup ##\n\n\n### Environment \n\nSee [here](docs/env_setup.md) for setting up the environment.\n\n\n### QM9\n\nThe dataset of QM9 will be automatically downloaded when running training.\n\n\n### MD17\n\nThe dataset of MD17 will be automatically downloaded when running training.\n\n\n### OC20\n\nThe dataset for different tasks can be downloaded by following instructions in their [GitHub repository](https://github.com/FAIR-Chem/fairchem/blob/b5a197fc3c79a9a5a787aabaa02979be53d296b7/DATASET.md#initial-structure-to-relaxed-structure-is2rs-and-initial-structure-to-relaxed-energy-is2re-tasks), which was previously this [link](https://github.com/Open-Catalyst-Project/ocp/blob/main/DATASET.md#download-and-preprocess-the-dataset).\nAfter downloading, place the datasets under `datasets/oc20/` by using `ln -s`.\nTake `is2re` as an example:\n\n```bash\n    cd datasets\n    mkdir oc20\n    cd oc20\n    ln -s ~/ocp/data/is2re is2re\n```\n\n\n\n## Training ##\n\n\n### QM9\n\n1. We provide training scripts under [`scripts/train/qm9/equiformer`](scripts/train/qm9/equiformer).\nFor example, we can train Equiformer for the task of `alpha` by running:\n\n    ```bash\n        sh scripts/train/qm9/equiformer/target@1.sh\n    ```\n\n2. The QM9 dataset will be downloaded automatically as we run training for the first time.\n\n3. The target number for different regression tasks can be found [here](https://pytorch-geometric.readthedocs.io/en/latest/generated/torch_geometric.datasets.QM9.html#torch_geometric.datasets.QM9).\n\n4. We also provide the code for training Equiformer with linear messages and dot product attention.\nTo train Equiformer with linear messages, replace `--model-name 'graph_attention_transformer_nonlinear_l2'` with `--model-name 'graph_attention_transformer_l2'` in training scripts.\n\n5. The training scripts for Equiformer with linear messages and dot product attention can be found in [`scripts/train/qm9/dp_equiformer`](scripts/train/qm9/dp_equiformer).\n\n6. Training logs of Equiformer can be found [here](https://www.dropbox.com/sh/9yuah73ss9na8tw/AACWZSoqddgeRiL4kFuVEn7Ra?dl=0).\n\n\n### MD17\n\n1. We provide training scripts under [`scripts/train/md17/equiformer`](scripts/train/md17/equiformer).\nFor example, we can train Equiformer for the molecule of `aspirin` by running:\n\n    ```bash\n        sh scripts/train/md17/equiformer/se_l2/target@aspirin.sh    # L_max = 2\n        sh scripts/train/md17/equiformer/se_l3/target@aspirin.sh    # L_max = 3\n    ```\n\n2. Training logs of Equiformer with $L_{max} = 2$ and $L_{max} = 3$ can be found [here](https://www.dropbox.com/sh/yriizguj1m7t891/AACh395HiqCC0G7Lv4YURlyra?dl=0) ($L_{max} = 2$) and [here](https://www.dropbox.com/sh/1xrpzd1op0kwcpf/AABFYmgq0ZIprDE82RAY7oPSa?dl=0) ($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. \n\n3. 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).\nFor example, we train Equiformer with DeNS for the molecule of `aspirin` by running:\n    ```bash\n        sh scripts/train/md17/equiformer_dens/se_l2/target@aspirin.sh    # L_max = 2\n        sh scripts/train/md17/equiformer_dens/se_l3/target@aspirin.sh    # L_max = 3\n    ```    \n\n4. The logs of training Equiformer with $L_{max} = 2$ and $L_{max} = 3$ with DeNS can be found [here](https://www.dropbox.com/scl/fo/5zdtct6an40dhvlirgl7b/AEHcheg7pLoNy7QvGbkJZ0k?rlkey=nlkyel9njusjhe9vpisnezmf7\u0026st=12cjzhpy\u0026dl=0) ($L_{max} = 2$) and [here](https://www.dropbox.com/scl/fo/53ctmrx3uqxlj4mcbj5oz/AKNU-WipNnPWBKIyAhB4VZk?rlkey=r7gzv1mic79ea1mif5oy0pww9\u0026st=4tla8k1g\u0026dl=0) ($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. \n\n\n### OC20\n\n1. We train Equiformer on **IS2RE data only** by running:\n\n    ```bash\n        sh scripts/train/oc20/is2re/graph_attention_transformer/l1_256_nonlinear_split@all_g@2.sh\n    ```\n\n    a. This requires 2 GPUs and results in energy MAE of around 0.5088 eV for the ID sub-split of the validation set.\n\n    b. Pretrained weights and training logs can be found [here](https://www.dropbox.com/sh/gfmzs0yxz8g54pi/AABcnBweIUajQTJsjvSdjxf3a?dl=0).\n\n\n2. We train Equiformer on **IS2RE data with IS2RS auxiliary task and Noisy Nodes data augmentation** by running:\n\n    ```bash\n        sh scripts/train/oc20/is2re/graph_attention_transformer/l1_256_blocks@18_nonlinear_aux_split@all_g@4.sh\n    ```\n\n    a. This requires 4 GPUs and results in energy MAE of around 0.4156 eV for the ID sub-split of the validation set.\n\n    b. Pretrained weights and training logs can be found [here](https://www.dropbox.com/sh/kijz22aulnyptql/AAD4s54XOkY9BSZ-ZzRoviWNa?dl=0).\n\n\n## File Structure ##\n\nWe have different files and models for QM9, MD17 and OC20.\n\n### General \n\n1. [`nets`](nets) includes code of different network architectures for QM9, MD17 and OC20.\n2. [`scripts`](scripts) includes scripts for training models for QM9, MD17 and OC20.\n \n\n### QM9\n\n1. [`main_qm9.py`](main_qm9.py) is the training code for QM9 dataset.\n\n\n### MD17\n\n1. [`main_md17.py`](main_md17.py) is the code for training and evaluation on MD17 dataset.\n2. [`main_md17_dens.py`](main_md17_dens.py) extends `main_md17.py` so that we can train Equiformer with DeNS.\n\n\n### OC20\n\nSome differences are made to support:\n- Removing weight decay for certain parameters specified by `no_weight_decay`. One example is [here](nets/graph_attention_transformer_oc20.py#L247).\n- Cosine learning rate.\n\n1. [`main_oc20.py`](main_oc20.py) is the code for training and evaluating.\n2. [`oc20/trainer`](oc20/trainer) contains the code for energy trainers.\n3. [`oc20/configs`](oc20/configs) contains the config files for IS2RE.\n\n\n## Acknowledgement ##\n\nOur implementation is based on [PyTorch](https://pytorch.org/), [PyG](https://pytorch-geometric.readthedocs.io/en/latest/index.html), [e3nn](https://github.com/e3nn/e3nn), [timm](https://github.com/huggingface/pytorch-image-models), [ocp](https://github.com/Open-Catalyst-Project/ocp), [SEGNN](https://github.com/RobDHess/Steerable-E3-GNN) and [TorchMD-NET](https://github.com/torchmd/torchmd-net).\n\n\n\n## Citation ##\n\nIf you use our code or method in your work, please consider citing the following:\n\n```bibtex\n@inproceedings{\n    liao2023equiformer,\n    title={Equiformer: Equivariant Graph Attention Transformer for 3D Atomistic Graphs},\n    author={Yi-Lun Liao and Tess Smidt},\n    booktitle={International Conference on Learning Representations},\n    year={2023},\n    url={https://openreview.net/forum?id=KwmPfARgOTD}\n}\n```\n\nIf DeNS is helpful to your work, please consider citing the following as well:\n```bibtex\n@article{\n    DeNS,\n    title={Generalizing Denoising to Non-Equilibrium Structures Improves Equivariant Force Fields}, \n    author={Yi-Lun Liao and Tess Smidt and Muhammed Shuaibi* and Abhishek Das*},\n    journal={arXiv preprint arXiv:2403.09549},\n    year={2024}\n}\n```\n\nPlease direct any questions to Yi-Lun Liao (ylliao@mit.edu).","project_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fatomicarchitects%2Fequiformer","html_url":"https://awesome.ecosyste.ms/projects/github.com%2Fatomicarchitects%2Fequiformer","lists_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fatomicarchitects%2Fequiformer/lists"}