AMS-ICAMS-RUB/grace-foundation-models

Pretrained GRACE Foundation Models

This repository distributes pretrained GRACE (Graph Atomic Cluster Expansion) machine-learning interatomic potentials, fitted with GRACEmaker / tensorpotential. The model tables and usage notes below are adapted from the GRACEmaker documentation.

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

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., arXiv

The base models (-OMAT) are trained on the OMat24 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., arXiv

These models are first pre-trained on OMat24 and then fine-tuned on a combination of the sAlex dataset (10.4M structures) and the MPtraj 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.

---

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:

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:

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:

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

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

Discover variable names with:

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