|
--- |
|
license: apache-2.0 |
|
library_name: transformers |
|
datasets: |
|
- vector-institute/s2ef-15m |
|
metrics: |
|
- mae |
|
pipeline_tag: graph-ml |
|
--- |
|
|
|
# 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: |
|
|
|
```python |
|
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](https://github.com/VectorInstitute/AtomGen) 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 | |
