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:
Task | SMP | PIP | RES | MSP | LBA | LEP | PSR | RSR |
---|---|---|---|---|---|---|---|---|
1.077 | TBD | TBD | TBD | TBD | TBD | TBD | TBD |
- Downloads last month
- 28