Open-source project
atomicarchitects/equiformer_v2 avatar
atomicarchitects/equiformer_v2

EquiformerV2: An Equivariant Transformer for OC20, and What Its Training Scripts Actually Require

[ICLR 2024] EquiformerV2: Improved Equivariant Transformer for Scaling to Higher-Degree Representations

351 stars49 forksPythonMIT

At a glance

What is it?
EquiformerV2 is the official PyTorch implementation of an ICLR 2024 equivariant transformer for atomistic energy and force prediction. Its value is concentrated in the OC20 training pipeline, and its cost is measured in multi-node GPU runs and a patched upstream dependency.
Who is it for?
Adopt EquiformerV2 if your task is OC20 S2EF energy and force prediction and you can supply the GPUs its training scripts assume, or if you only need to run the published 2M and All+MD checkpoints for evaluation. Do not adopt it as a general-purpose molecular property model or as a small-scale experiment: the repository ships OC20 S2EF code only, the README points elsewhere for environment setup, and the multi-node scripts assume clusters of 8 to 16 nodes.
Can I use it commercially?
Yes. MIT is a permissive licence: you can use, modify and sell software built on it, as long as you keep its copyright and licence notices.
Is it still maintained?
Yes. The repository last received commits 159 days ago.
What is it written in?
Mainly Python, according to GitHub's language statistics.

Answers come from the project's GitHub data, last synced on September 15, 2026, and from our analysis. They are not legal advice.

DEEP OPEN-SOURCE ANALYSIS

The problem EquiformerV2 addresses: energy and force prediction on OC20

The repository targets one benchmark family. The README states that it provides code for training the base model setting on the OC20 S2EF-2M and S2EF-All+MD datasets. S2EF means Structure to Energy and Forces: given an atomic structure, predict a total energy and a per-atom force vector. That is the substrate for interatomic potentials and for the relaxation and catalyst screening work the OC20 project is built around, which is why the repository topics include interatomic-potentials, force-fields, catalyst-design and molecular-dynamics.

The intended user is a researcher or engineer who already works inside the Open Catalyst ecosystem, not someone looking for a drop-in molecular property predictor. The README points at the OCP repository for dataset download and at a separate docs/env_setup.md for the environment, which tells you the project assumes familiarity with that stack. The paper title frames the contribution as scaling to higher-degree representations, meaning the model's internal representation of geometry is expanded beyond what earlier equivariant transformers used. Whether that scaling pays off is a question the paper answers; the repository's job is to let you reproduce it.

What the repository layout tells you about the architecture

The file structure is small and legible. nets holds network architectures for OC20, scripts holds training scripts, main_oc20.py is the entry point for training, evaluating and running relaxation, oc20/trainer contains the force trainer and utility functions, and oc20/configs contains the S2EF configuration files.

That split matters when you are deciding whether to modify the model. The architecture lives in nets, the training loop and force loss handling live in oc20/trainer, and the hyperparameters you would most likely touch live in the YAML files under oc20/configs. The README's checkpoint table links each checkpoint to a matching config, for example oc20/configs/s2ef/2M/equiformer_v2/equiformer_v2_N@12_L@6_M@2_epochs@30.yml. Those filenames encode the model shape: N is the number of channels, L the number of layers, M the maximum degree. The 2M checkpoint is N@12 L@6 M@2; the 31M parameter model is N@8 L@4 M@2; the 153M model is N@20 L@6 M@3. So the scaling claim in the title is visible directly in the naming convention: the largest model raises M from 2 to 3 and widens N from 12 to 20.

The README also notes that EquiformerV2 has been incorporated into the OCP repository and used in the Open Catalyst demo. If you want the model as a component inside a maintained framework rather than as a standalone training repo, that is the more direct route.

Getting the environment and data in place

The README does not inline the environment instructions. It says see docs/env_setup.md for setting up the environment, and that file is not reproduced here, so the exact dependency versions are outside what can be confirmed from this material. What is confirmed is that the project is Python and PyTorch, and that it depends on the ocp codebase.

Dataset preparation is explicit. The OC20 S2EF dataset is downloaded by following the Open Catalyst Project instructions, and the README gives this example for the 2M split:

cd ocp python scripts/download_data.py --task s2ef --split "2M" --num-workers 8 --ref-energy

The README adds that you also need the val_id data split to run training. After downloading, the datasets are placed under datasets/oc20/ using symlinks:

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

One warning deserves attention before you spend time debugging. The README states that if there is an import error, it is possible that ocp/ocpmodels/common/utils.py is not modified, and directs you to docs/env_setup.md. In other words, this repository requires a local patch to an upstream dependency. That is a real maintenance cost: the patch is not expressed as a pinned fork or a patch file in the material shown, so anyone reproducing the setup has to read the environment doc and apply it by hand.

The training commands and the hardware they assume

Training is launched through shell scripts, not a Python CLI. For the 2M split the README gives:

sh scripts/train/oc20/s2ef/equiformer_v2/equiformer_v2_N@12_L@6_M@2_splits@2M_g@multi-nodes.sh

and notes that this script uses 2 nodes with 8 GPUs on each node. A single-node variant exists:

sh scripts/train/oc20/s2ef/equiformer_v2/equiformer_v2_N@12_L@6_M@2_splits@2M_g@8.sh

for 8 GPUs on 1 node. The All+MD runs are larger. The 153M model uses equiformer_v2_N@20_L@6_M@3_splits@all+md_g@multi-nodes.sh across 16 nodes with 8 GPUs each, and the 31M model uses equiformer_v2_N@8_L@4_M@2_splits@all+md_g@multi-nodes.sh across 8 nodes with 8 GPUs each.

Read those numbers as a gate. The full All+MD configuration assumes 128 GPUs. The smallest published configuration still assumes 8. If your available hardware is a single workstation GPU, this repository is not the path of least resistance; the released checkpoints are, and the README provides them. The scripts also encode the distributed launch configuration, so adapting them to a different topology means editing shell scripts rather than passing flags to a documented interface.

Checkpoints, and what the published numbers do and do not tell you

Three checkpoints are offered. EquiformerV2 on the 2M split reports val force MAE of 19.4 meV/Å and val energy MAE of 278 meV. The 31M model on All+MD reports 16.3 and 232. The 153M model on All+MD reports 15.0 and 227. Each row links both a checkpoint file and a config, so you can inspect the exact settings that produced the number.

Two cautions. First, the 2M checkpoint is described as trained for 30 epochs, and its config filename carries epochs@30, so the training budget is part of the definition of that result. Second, these are validation MAEs on OC20 splits. They say nothing about transfer to a different chemistry, a different level of theory, or a system size outside the benchmark distribution. The README does not make claims beyond these splits, and neither should a deployment decision.

The README also points to a follow-up: DeNS, denoising non-equilibrium structures, a self-supervised objective analogous to BERT applied to 3D atomistic systems, reported to improve energy and force prediction for EquiformerV2. That work lives in a separate repository, so adopting it is a second integration, not a configuration flag here.

Where EquiformerV2 is the wrong tool

The clearest limitation is scope. The README describes code for the base model setting on two OC20 S2EF datasets. There is no material here describing inference APIs, pretrained models for other datasets, or support for molecular property prediction tasks generally. The topic list is broad, but the shipped code is narrow.

A second limitation is the dependency patch. Requiring a modification to an upstream file means your environment diverges from a stock ocp install. If you later upgrade ocp, you have to re-check whether the modification still applies and whether the surrounding code changed. The README flags the symptom (an import error) but the fix lives in a separate document.

A third consideration is that the project has a named successor. The README states that the next generation, EquiformerV3, is available at a separate repository, with further improvements to efficiency, expressivity and generality, and results on tasks requiring higher-order derivatives. For a new project starting today, that is the more relevant starting point unless you specifically need to reproduce the ICLR 2024 results. The README does not present EquiformerV2 as unmaintained, and the repository is not archived, but the authors themselves position it as a prior generation.

Alternatives and the difference in approach

The most direct alternative is the OCP repository itself. The README states that EquiformerV2 has been incorporated into OCP at src/fairchem/core/models/equiformer_v2 and used in the Open Catalyst demo. The difference is packaging rather than mathematics: OCP gives you the same architecture inside a framework that also handles dataset loading, training infrastructure and evaluation for the broader Open Catalyst task family, whereas this repository gives you the model plus a set of OC20-specific training scripts and configs. If your goal is to train or fine-tune on OC20 with less environment work, OCP is the shorter route. If your goal is to read and modify the architecture itself, this repository is the more direct view of it, because nets and oc20/trainer are exposed without the surrounding framework.

A second alternative is EquiformerV3, in the same authors' organization. The README frames the difference in terms of efficiency, expressivity and generality, and specifically mentions tasks requiring higher-order derivatives such as gradients of forces. That is a capability distinction, not a cosmetic one: if your downstream work needs second derivatives, the README's own description points you to V3.

Licence and maintenance expectations

The repository is MIT licensed. For practical purposes that is a permissive licence, and it is the same licence family used by much of the surrounding scientific Python ecosystem, but the material here does not include the full licence text or any statement about third-party components. The code depends on the ocp codebase and on eSCN, which the README cites as a separate paper; their licences are separate from this one and are not described in the material provided. Check those before redistributing a combined artifact. This is a description of what the repository states, not legal advice.

On maintenance, the material shows no releases retrieved, so there is no versioned artifact to pin against. Updates arrive as commits to main. The changelog lives at docs/changelog.md and is not reproduced here, so the rate and nature of changes cannot be assessed from this material. The practical consequence is that reproducing a result means recording the commit you used, since there is no release tag to point at. Upgrading later means diffing against main and re-verifying the ocp/utils.py modification, which is the piece most likely to break when the upstream dependency moves.

Editorial conclusion

Adopt EquiformerV2 if your task is OC20 S2EF energy and force prediction and you can supply the GPUs its training scripts assume, or if you only need to run the published 2M and All+MD checkpoints for evaluation. Do not adopt it as a general-purpose molecular property model or as a small-scale experiment: the repository ships OC20 S2EF code only, the README points elsewhere for environment setup, and the multi-node scripts assume clusters of 8 to 16 nodes. Verify three things before committing: that you have applied the documented modification to ocp/ocpmodels/common/utils.py, that your OC20 download includes the val_id split the training runs need, and whether you want the archived ICLR 2024 implementation or the successor EquiformerV3 that the README names as the next generation.

Official sources

  1. atomicarchitects/equiformer_v2 on GitHub
  2. Issues
  3. License: MIT
  4. Project website
  5. README
Community notes

Community notes