mace_omat_medium — Fine-tuned MACE OMAT for LiF

A LoRA fine-tuned MACE (Multi-Atomic Cluster Expansion) foundation model for ionic lithium–fluoride systems. This model is built on the MACE-MP-0 OMAT medium checkpoint and fine-tuned on Quantum ESPRESSO DFT data covering bulk LiF, LiF interfaces, isolated frames.

This model is directly produced by the mlpdft project.

Model Details

Model Description

MACE (Multi-Atomic Cluster Expansion) is an equivariant message-passing neural network that constructs many-body expansions with learnable radial and angular features. It achieves state-of-the-art accuracy while maintaining speed suitable for molecular dynamics simulations of thousands of atoms.

This variant is fine-tuned via LoRA (Low-Rank Adaptation) on the MACE-MP-0 OMAT-medium foundation model to improve transfer to ionic lithium–fluoride chemistries.

Developed by: Jorge Munoz (mlpdft project)
Base model: MACE-MP-0 OMAT medium (mace_omat_medium)
Model type: Equivariant message-passing neural network (E(3)-equivariant)
Fine-tuning method: LoRA (Low-Rank Adaptation)
License: MIT

Model Sources

Repository: github.com/jorgemunozl/mlpdft
Base model: MACE-MP-0
Paper (MACE foundation): A foundation model for atomistic materials chemistry

Training Details

Training Data

The model was fine-tuned on the jorgemunozl/minimal_li_f_mace_dataset dataset, which contains DFT calculations (Quantum ESPRESSO, PBEsol, PAW pseudopotentials) spanning 8 groups:

Group	Description
`LIF64_KJPAW_V2`	Bulk LiF — NVE / NPT trajectories
`LIF64_ISOLATED`	Isolated bulk LiF frames
`LIFINTERFACE_KJPAW_V1`	LiF interface (first version)
`LIFINTERFACE_KJPAW_NPT`	LiF interface — NPT
`LIFINTERFACE_KJPAW_NPT_V2`	LiF interface — NPT (second version)
`LIWITHF_V3`	Li + F slabs
`LIWITHF_ISOLATED`	Isolated Li + F frames
`LIWITHF_NPT_FINAL`	Li + F — NPT (final)

Each frame carries:

REF_energy — total DFT energy (eV)
REF_forces — per-atom forces (eV/Å)
stress — stress tensor (when available)

Energy offsets (isolated-atom energies in eV):

Element	Atomic number	Isolated atom energy
Li	3	`-15.11995 × RY_TO_EV`
F	9	`-58.46236 × RY_TO_EV`

RY_TO_EV ≈ 13.605698

Training Procedure

Fine-tuning hyperparameters

Hyperparameter	Value
Architecture	MACE
Cutoff radius (`r_max`)	5.0 Å
Message-passing channels	128
Max spherical harmonics (`max_L`)	1 (scalars + vectors)
Max ell order (`max_ell`)	3
Interaction blocks	2
Body order (correlation)	3 (4-body)
Radial basis functions	8
Cutoff basis functions	5

Optimization
Optimizer	Adam (AMSGrad)
Learning rate	0.01
Weight decay	5 × 10⁻⁷
Gradient clipping	10.0
Scheduler	ReduceLROnPlateau
LR decay factor	0.8
Scheduler patience	50 epochs
Max epochs	10
Batch size	8
Early stopping patience	4
Validation fraction	50%

Loss	Weighted
Energy weight	1.0
Forces weight	1.0

LoRA fine-tuning
LoRA enabled	Yes
LoRA rank	8

Regularization
SWA (stochastic weight averaging)	Disabled
EMA (exponential moving average)	Yes (decay 0.99)

Precision	float64

Training script

Fine-tuning is done via the mace_run_train CLI from the mace-torch package, driven by src/mlpdft/train.py. See the mlpdft README for full details.

uv run python src/mlpdft/train.py

Hardware

Training was performed on CUDA GPU RTX4000ADA support (device="cuda").

Evaluation

Metrics

Evaluation metrics are computed per group on energy and forces:

Metric	Description
MAE	Mean Absolute Error
RMSE	Root Mean Square Error
MaxAE	Maximum Absolute Error
MAE per atom	Energy MAE divided by number of atoms

Evaluation is run via:

uv run python src/mlpdft/evaluate_mace_metrics.py

The script runs the model on all groups from the dataset and prints a summary table with energy and force errors.

Technical Specifications

Model Architecture and Objective

MACE is a body-ordered equivariant message-passing neural network:

Input: Atomic numbers and Cartesian positions
Embeddings: One-hot atomic number → learnable scalar features
Message passing: Atomic cluster expansion (ACE) basis with radial Bessel basis + spherical harmonics for angular information
Equivariance: E(3)-equivariant (rotations, translations, reflections)
Readout: Site energies summed to total energy; forces via automatic differentiation
Output: Total energy (eV) and per-atom forces (eV/Å)

Compute Infrastructure

Framework: PyTorch 2.x + mace-torch ≥ 0.3.6
Hardware requirements: CPU for inference; GPU recommended for large-scale MD

Citation

If you use this model, please cite both the MACE foundation model and the CHGNet paper:

@article{batatia2023foundation,
      title={A foundation model for atomistic materials chemistry},
      author={Ilyes Batatia and Philipp Benner and Yuan Chiang and Alin M. Elena
              and Fabian Zills and Gábor Csányi},
      year={2023},
      eprint={2401.00096},
      archivePrefix={arXiv},
      primaryClass={physics.chem-ph}
}

@article{deng2023chgnet,
      title={CHGNet: Pretrained universal neural network potential
             for charge-informed atomistic modeling},
      author={Bowen Deng and Peichen Zhong and KyuJung Jun
              and Janosh Riebesell and Kevin Han and Christopher J. Bartel
              and Gerbrand Ceder},
      year={2023},
      eprint={2302.14231},
      archivePrefix={arXiv},
      primaryClass={cond-mat.mtrl-sci}
}

MACE-Universal by Yuan Chiang, 2023, Hugging Face, Revision e5ebd9b, DOI: 10.57967/hf/1202, URL: https://huggingface.co/cyrusyc/mace-universal