vector-institute / atomformer-base

AtomFormer base model

This model is a transformer-based model that leverages gaussian pair-wise positional embeddings to train on atomistic graph data. It is part of a suite of datasets/models/utilities in the AtomGen project that supports other methods for pre-training and fine-tuning models on atomistic graphs.

Model description

AtomFormer is a transformer model with modifcations to train on atomstic graphs. It builds primarily on the work from uni-mol+ to add the pair-wise pos. embeds. to the attention mask to leverage 3-D positional information. This model was pre-trained on a diverse set of aggregated atomistic datasets where the target task is the per-atom force prediction and the per-system energy prediction.

The model also includes metadata regarding the atomic species that are being modeled, this includes the atomic radius, electronegativity, valency, etc. The metadata is normalized and projected to be added to the atom embeddings in the model.

Intended uses & limitations

You can use the raw model for either force and energy prediction, but it's mostly intended to be fine-tuned on a downstream task. The performance of the model as a force and energy prediction model is not validated, it was primarily used a pre-training task.

How to use

Here is how to use the model to extract features from the pre-trained backbone:

import torch
from transformers import AutoModel
model = AutoModel.from_pretrained("vector-institute/atomformer-base",
                                   trust_remote_code=True)

input_ids = torch.randint(0, 50, (1, 10))
coords = torch.randn(1, 10, 3)
attn_mask = torch.ones(1, 10)

output = model(input_ids, coords=coords, attention_mask=attention_mask)
output.shape #  (torch.Size([1, 10, 768])

Training data

AtomFormer is trained on an aggregated S2EF dataset from multiple sources such as OC20, OC22, ODAC23, MPtrj, and SPICE with structures and energies/forces for pre-training. The pre-training data includes total energies and formation energies but trains using formation energy (which isn't included for OC22, indicated by "has_formation_energy" column).

Preprocessing

The model expects input in the form of tokenized atomic symbols represented as input_ids and 3D coordinates represented as coords . For the pre-training task it also expects labels for the forces and formation_energy .

The DataCollatorForAtomModeling utility in the AtomGen library has the capacity to perform dynamic padding to batch the data together. It also offers the option to flatten the data and provide a batch column for gnn-style training.

Pretraining

The model was trained on a node of 4xA40 (48 GB) for 10 epochs (~2 weeks). See the training code for all hyperparameters details.

Evaluation results

We use the Atom3D dataset to evaluate the model's performance on downstream tasks.

When fine-tuned on downstream tasks, this model achieves the following results: