v0.3.14

[ilyes319]

MACE v0.3.14 Release Notes

We are excited to announce MACE v0.3.14, featuring significant new capabilities for embedding functionality, GPU acceleration, dielectric properties prediction, and enhanced training options.

🏗️ Foundation Models

MACE-OMOL Foundation Models

Introduced support for MACE-OMOL models trained on the 100M OMOL dataset with charge and spin embeddings, providing improved accuracy for organic molecules with charges and spins.

Example usage:

from mace.calculators import mace_omol

# Load MACE-OMOL model with charge and spin support
calc = mace_omol(model="extra_large", device="cuda")
atoms.calc = calc

# Set charge and spin for the system
atoms.info["charge"] = 1.0      # +1 charge
atoms.info["spin"] = 1.0        # spin multiplicity

energy = atoms.get_potential_energy()
forces = atoms.get_forces()

Small OMAT Model

Added omat-small-0 to foundation model shortcuts for faster inference when high accuracy is not critical.

Example usage:

from mace.calculators import mace_mp
calc = mace_mp(model="small-omat-0")  # Faster, smaller model

🎯 Fine-tuning

Pseudo-label Finetuning

Added multihead pseudo-label finetuning capability, allowing models to be fine-tuned on their own predictions for improved accuracy and easy of use.

Example usage:

python run_train.py \
  --foundation_model="medium" \
  --train_file=real_data.xyz \
  --pt_train_file=replay_data.xyz \
  --multiheads_finetuning=True \
  --pseudolabel_replay=True \

Enhanced Multihead Finetuning

Fixed various issues with multihead finetuning functionality and improved robustness for the code.
Add option to use any of the models from mace_mp as --foundation_model key name, see:

mace/mace/calculators/foundations_models.py

Line 19 in 0139da1

mace_mp_urls = {

. For example, --foundation_model="small-omat-0".

🔬 Models

Embedding Functionality for MACE

Added support for embedding additional properties like total charge, total spin, and electronic temperature into MACE models, enabling more accurate predictions for charged and magnetic systems.

Training example with custom embeddings:

python run_train.py \
  --train_file=data.xyz \
  --embedding_specs='{"charge": {"embed_type": "continuous", "min": -5, "max": 5}, "spin": {"embed_type": "continuous", "min": 0, "max": 4}, "temperature": {"embed_type": "continuous", "min": 0, "max": 3000}}' \
  --use_embedding_readout \
  --model=MACE

See code and test.

Dielectric MACE for Polarizability Prediction

Added support for the Dielectric MACE model to predict molecular polarizabilities alongside energies and forces.

Reference: Kapil, et al, "First-principles spectroscopy of aqueous interfaces using machine-learned electronic and quantum nuclear effects"

Example usage:

# Train dielectric MACE model
python run_train.py \
  --model="AtomicDielectricMACE" \
  --train_file=data.xyz \
  --loss="dipole_polar" \
  --polarizability_key="REF_polarizability" \
  --error_table="DipolePolarRMSE" \

See code and test.

LES Models for Coulomb Interactions

Implemented support for LES models for improved description of Coulomb interactions.

Reference: Bingqing Cheng, "Latent Ewald summation for machine learning of long-range interactions"

Example usage:

# Train MACELES model
python run_train.py \
  --model=MACELES \
  --train_file=data.xyz \

See code and test.

New Non-linear Blocks

Introduced improved non-linear readout blocks that provide better accuracy when training MACE on large datasets. See code.

⚡ Performance Improvements

OpenEquivariance Support

Integrated OpenEquivariance support for MACE kernels, enabling acceleration on AMD GPUs and providing alternative compute backends. Use --enable_oeq=True for training with openequivariance, and use enable_oeq=True to evaluate ASE calc with openequivariance.

CuEq Fused Operations

Added CuEq fused operations for tensor product and scatter sum operations, significantly improving training and inference speed on CUDA devices.

Refactored Clebsch-Gordan Coefficients

Implemented option to use refactored CG coefficients with fewer computational paths for improved efficiency.

Example usage:

python run_train.py \
  --train_file=data.xyz \
  --use_reduced_cg

Element Agnostic Product Basis

Added option for element-agnostic product basis, reducing memory usage and improving scalability for systems with many element types.

Example usage:

python run_train.py \
  --train_file=data.xyz \
  --use_agnostic_product

🔧 Training and Infrastructure Improvements

Improved Multi-GPU Support

• Fixed support for torchrun and MPI multi-GPU training
• Dynamically set backend for distributed training based on device type
• Better handling of distributed training configurations

Intel GPU Support

Added IPEX (Intel Extension for PyTorch) support for Intel GPU inference and training.

🐛 Bug Fixes and Improvements

• Fixed bug preventing building 1-layer (no message passing) MACE models
• Improved handling of E0s specification via JSON in preprocessing
• Better pathlib usage for file suffix handling in finetuning_select
• Fixed OEQ import handling on macOS without CUDA
• Corrected loss output to log files during fine-tuning
• Improved behavior of auto-downloaded MP data when using pseudolabels
• Enhanced handling of atomic numbers in multihead training

📚 Additional Features

• Added support for calculating descriptors in eval_configs
• Enhanced node energy output capabilities
• Improved interaction energy plotting functionality
• Better filtering of training quantities with zero weights

🙏 Acknowledgments

We thank all contributors to this release, including new contributors @Alexsp32, @Enry99, @vbharadwaj-bk, @naik-aakash, and @Nilsgoe.

Full Changelog: v0.3.13...v0.3.14

---

For detailed documentation and examples, visit our GitHub repository and documentation.

---

This discussion was created from the release v0.3.14.