README.md

---
license: other
license_name: academic-software-license
license_link: https://github.com/ICAMS/grace-tensorpotential/blob/master/LICENSE.md
library_name: tensorpotential
tags:
  - grace
  - interatomic-potential
  - molecular-dynamics
  - machine-learning-potential
  - materials-science
---

# Pretrained GRACE Foundation Models

This repository distributes pretrained **GRACE** (Graph Atomic Cluster Expansion) machine-learning
interatomic potentials, fitted with [GRACEmaker / tensorpotential](https://github.com/ICAMS/grace-tensorpotential).
The model tables and usage notes below are adapted from the
[GRACEmaker documentation](https://github.com/ICAMS/grace-tensorpotential/blob/master/docs/gracemaker/foundation.md).

> **Use the "Full Name" column** to refer to a model in LAMMPS and ASE.

The **UQ** and **Kokkos** columns indicate, per model:

- **UQ** — the distributed SavedModel ships with an uncertainty-quantification head
  (and the checkpoint includes `gmm_artifacts.npz`).
- **Kokkos** — the model archive bundles `kokkos.npz` for LAMMPS-Kokkos deployment.

## License

**All GRACE models distributed in this repository are released under the Academic Software
License (ASL).** By downloading or using any model from this repository you agree to the terms
of the ASL: <https://github.com/ICAMS/grace-tensorpotential/blob/master/LICENSE.md>

---

## SMAX models

Reference: [arXiv](https://arxiv.org/abs/2602.23489)

The **SMAX** (**Maximum Entropy**) models are trained on a chemistry-agnostic dataset generated via a multicomponent
maximum information entropy structure generation protocol.

Unlike traditional datasets that focus on low-energy equilibrium structures, SMAX is constructed to deliberately sample
broad and diverse regions of configurational space. This provides a robust physical prior for atomic interactions across the entire periodic table, enabling accurate modeling of large-strain phase transformations, defects in complex alloys, and reaction barriers in catalytic systems.

**Custom Cutoffs & Interaction Ranges:** SMAX models utilize a **custom element-dependent cutoff radius** ranging from
**5.0 Å to 7.5 Å**.

**Recommendation:** For most general-purpose applications, we recommend using the **SMAX-OMAT** models.
They offer the best balance of structural robustness (from SMAX) and high-precision energy/force accuracy (from OMat24).

#### Single-layer, local models

| Model Name | Full Name | Size | κ_SRME | UQ | Kokkos | Description |
| :---- | :---- | :---- | :---- | :---- | :---- | :---- |
| GRACE-1L-SMAX-L | GRACE-1L-SMAX-large | large | 0.696 | — | ✓ | Single-layer local (SMAX) |
| **GRACE-1L-SMAX-OMAT-L** | **GRACE-1L-SMAX-OMAT-large** | **large** | 0.338 | — | ✓ | **Single-layer local (SMAX + OMat24)** |

#### Two-layer, semilocal models

| Model Name | Full Name | Size | κ_SRME | UQ | Kokkos | Description |
| :---- | :---- | :---- | :---- | :---- | :---- | :---- |
| GRACE-2L-SMAX-M | GRACE-2L-SMAX-medium | medium | 0.469 | — | ✓ | Two-layer semi-local (SMAX) |
| GRACE-2L-SMAX-L | GRACE-2L-SMAX-large | large | 0.444 | — | ✓ | Two-layer semi-local (SMAX) |
| **GRACE-2L-SMAX-OMAT-M** | **GRACE-2L-SMAX-OMAT-medium** | **medium** | 0.197 | — | ✓ | **Two-layer semi-local (SMAX + OMat24)** |
| **GRACE-2L-SMAX-OMAT-L** | **GRACE-2L-SMAX-OMAT-large** | **large** | 0.191 | — | ✓ | **Two-layer semi-local (SMAX + OMat24)** |

## OMAT models

Reference: [npj Comp. Mat.](https://www.nature.com/articles/s41524-026-01979-1), [arXiv](https://arxiv.org/abs/2508.17936)

The base models (**-OMAT**) are trained on the [OMat24](https://huggingface.co/datasets/fairchem/OMAT24#omat24-dataset) dataset.
The fine-tuned versions (**-OMAT-ft-E**) are derived from these base models by fine-tuning with more emphasis on energies.
All models listed use a fixed **6 Å cutoff**.

#### Single-layer, local models

| Model Name | Full Name | Size | κ_SRME | UQ | Kokkos | Description |
|:---|:---|:---|:---|:---|:---|:---|
| **GRACE-1L-OMAT** | GRACE-1L-OMAT | small | 0.398 | — | — | Single-layer local |
| GRACE-1L-OMAT-M-base | GRACE-1L-OMAT-medium-base | medium | 0.380 | ✓ | ✓ | Single-layer local (base) |
| **GRACE-1L-OMAT-M** | GRACE-1L-OMAT-medium-ft-E | medium | 0.417 | ✓ | ✓ | Single-layer local (finetuned on energy) |
| GRACE-1L-OMAT-L-base | GRACE-1L-OMAT-large-base | large | **0.354** | ✓ | ✓ | Single-layer local (base) |
| **GRACE-1L-OMAT-L** | GRACE-1L-OMAT-large-ft-E | large | 0.383 | ✓ | ✓ | Single-layer local (finetuned on energy) |

#### Two-layer, semilocal models

| Model Name | Full Name | Size | κ_SRME | UQ | Kokkos | Description |
|:---|:---|:---|:---|:---|:---|:---|
| **GRACE-2L-OMAT** | GRACE-2L-OMAT | small | 0.288 | — | — | Two-layer semi-local |
| GRACE-2L-OMAT-M-base | GRACE-2L-OMAT-medium-base | medium | 0.212 | ✓ | ✓ | Two-layer semi-local (base) |
| **GRACE-2L-OMAT-M** | GRACE-2L-OMAT-medium-ft-E | medium | 0.217 | ✓ | ✓ | Two-layer semi-local (finetuned on energy) |
| GRACE-2L-OMAT-L-base | GRACE-2L-OMAT-large-base | large | **0.165** | ✓ | ✓ | Two-layer semi-local (base) |
| **GRACE-2L-OMAT-L** | GRACE-2L-OMAT-large-ft-E | large | 0.186 | ✓ | ✓ | Two-layer semi-local (finetuned on energy) |

## OAM models

Reference: [npj Comp. Mat.](https://www.nature.com/articles/s41524-026-01979-1), [arXiv](https://arxiv.org/abs/2508.17936)

These models are first pre-trained on **OMat24** and then fine-tuned on a combination of the [sAlex](https://huggingface.co/datasets/fairchem/OMAT24#salex-dataset) dataset (10.4M structures) and the [MPtraj](https://figshare.com/articles/dataset/Materials_Project_Trjectory_MPtrj_Dataset/23713842?file=41619375) dataset (1.58M structures).

#### Single-layer, local models

| Model Name | Full Name | Size | F1 | κ_SRME | UQ | Kokkos | Description |
| :--- |:---| :--- |:---|:---|:---|:---| :--- |
| GRACE-1L-OAM | GRACE-1L-OAM | small | 0.824 | 0.516 | — | — | Single-layer local |
| GRACE-1L-OAM-M | GRACE-1L-OMAT-medium-ft-AM | medium | 0.800 | 0.411 | — | ✓ | Single-layer local |
| GRACE-1L-OAM-L | GRACE-1L-OMAT-large-ft-AM | large | 0.815 | **0.377** | — | ✓ | Single-layer local |

#### Two-layer, semilocal models

| Model Name | Full Name | Size | F1 | κ_SRME | UQ | Kokkos | Description |
| :--- |:--- | :--- |:---|:---|:---|:---| :--- |
| GRACE-2L-OAM | GRACE-2L-OAM | small | 0.880 | 0.294 | — | — | Two-layer semi-local |
| GRACE-2L-OAM-M | GRACE-2L-OMAT-medium-ft-AM | medium | 0.881 | 0.200 | — | ✓ | Two-layer semi-local |
| GRACE-2L-OAM-L | GRACE-2L-OMAT-large-ft-AM | large | 0.889 | **0.168** | — | ✓ | Two-layer semi-local |

## Mixed-precision (-mx) variants

Mixed-precision builds of the 2L-large models: float32 weights, which the Kokkos export
promotes to float64. Energies and forces agree with the TF SavedModel to ~10⁻⁵ (vs ~10⁻⁷ for
the full-precision models) — the expected effect of fp32 accumulation in the SavedModel.

| Full Name | Size | UQ | Kokkos | Description |
|:---|:---|:---|:---|:---|
| GRACE-2L-OMAT-large-mx | large | — | ✓ | Two-layer semi-local, mixed precision (OMat24) |
| GRACE-2L-OMAT-large-mx-ft-AM | large | — | ✓ | Two-layer semi-local, mixed precision (OMat24, fine-tuned on sAlex + MPtraj) |

---

## Repository layout

Each model is provided as two `.tar.gz` archives:

- `models/<full-name>-model.tar.gz` — the deployable TensorFlow `SavedModel` directory
  (used by the ASE `TPCalculator`). UQ-enabled models embed an uncertainty head, and many
  archives also bundle `kokkos.npz` for LAMMPS-Kokkos deployment.
- `checkpoints/<full-name>-checkpoint.tar.gz` — the training checkpoint required for
  fine-tuning (`model.yaml`, `checkpoint.index`, `checkpoint.data-*`, plus `gmm_artifacts.npz`
  for UQ models).

## Downloading Foundation Models

The recommended way is the `grace_models` utility shipped with `tensorpotential`:

- `grace_models list` — view the list of available models.
- `grace_models download <model_name>` — download a specific model (SavedModel).
- `grace_models checkpoint <model_name>` — download a model's training checkpoint.

By default all foundation models are stored in `$HOME/.cache/grace`; override with the
`GRACE_CACHE` environment variable. The downloaded models can be used for simulations in
**ASE** and **LAMMPS** — see the
[GRACEmaker quickstart](https://github.com/ICAMS/grace-tensorpotential/blob/master/docs/quickstart.md).

---

## Fine-tuning foundation models

Fine-tuning is performed from **checkpoints** (not SavedModels). Run `grace_models list`
to see which models expose a `CHECKPOINT:` field; download with
`grace_models checkpoint <MODEL-NAME>` (or it is fetched automatically when needed).

Generate a fine-tuning `input.yaml` interactively with `gracemaker -t`, or add to the
`potential` section manually:

```yaml
potential:
  finetune_foundation_model: GRACE-1L-OAM
  shift: auto  # automatically align FM energies with your DFT reference data
# reduce_elements: True # if True - select from original models only elements present in the CURRENT dataset
```

### Automatic energy shift correction (`shift: auto`)

When `shift: auto` is set alongside `finetune_foundation_model`, GRACEmaker computes optimal
per-element energy shifts that minimize the difference between the foundation model's
predictions and your reference DFT data, injecting them via `ConstantScaleShiftTarget` before
training. This is useful when your DFT settings differ from the foundation model's training
data (systematic energy offset) and leads to faster convergence.

`shift: auto` is independent of `data::reference_energy`. Typical usage:

* `reference_energy: 0` (or dict) — set your DFT reference frame
* `shift: auto` — align the FM output to your DFT reference frame

### Learning rate & initial metrics

Use a small `learning_rate` and evaluate initial metrics:

```yaml
fit:
  opt_params: {learning_rate: 0.001, ... }
  eval_init_stats: True
```

### Frozen-weights fine-tuning

To mitigate catastrophic forgetting, restrict training to selected variables via
`trainable_variable_names`:

```yaml
# GRACE-2L models
fit:
  trainable_variable_names: ["I2/reducing_", "rho/reducing_", "I1/reducing_"]
```

```yaml
# GRACE-1L models
fit:
  trainable_variable_names: ["rho/reducing_"]
```

Discover variable names with:

```bash
grace_utils -p ~/.cache/grace/checkpoints/SOME_FOUNDATIONAL_MODEL_NAME/model.yaml summary -v 1
```