RussellKwon/UMA_s1p1_oc20_fine-tuned_WSO

UMA-S 1.1 Fine-Tuned Model for W-S-O Chemistry

A machine-learned interatomic potential (MLIP) for tungsten disulfide (WS₂) under oxygen ion (O⁺ and O₂⁺) bombardment, fine-tuned from Meta FAIR's UMA-S 1.1 backbone (OC20 task head). This is the production model used in our W-S-O simulations and in the accompanying paper, arXiv:2606.21632.

⚠️ License & provenance. This model is a derivative work of Meta's UMA-S 1.1 and its weights embed the UMA backbone. It is therefore distributed under the FAIR Chemistry License v1 (see LICENSE), not a generic open-source license. The base model facebook/UMA is gated. If you redistribute this model or any derivative, you must pass along the same FAIR Chemistry License (§1.b.i) and acknowledge the use of UMA in any publication (§1.b.ii).

Authors

Author	ORCID	Affiliation
Jaehong Kwon	0000-0002-2792-2772	Department of Chemical and Biological Engineering, Princeton University, Princeton, New Jersey 08544, USA
Andrew S. Rosen	0000-0002-0141-7006	Department of Chemical and Biological Engineering, Princeton University, Princeton, New Jersey 08544, USA
David B. Graves	0000-0002-2288-6768	Department of Chemical and Biological Engineering, Princeton University, Princeton, New Jersey 08544, USA; Princeton Plasma Physics Laboratory, Princeton, New Jersey 08540, USA

Corresponding author: David B. Graves (dgraves@princeton.edu).

Model details

Base model	uma-s-1p1 (facebook/UMA)
Task head	oc20
Fine-tuning method	Layer freezing (layers 2 and 4 frozen)
Trainable parameters	73,434,657 (50.1%); 73,131,520 frozen (49.9%)
Training targets	Energy, forces, and stress
Forces	Conservative (energy-gradient)

Training

• Data: 2,610 DFT-labeled configurations, fine-tuned with a 90/10 train/validation split. Labels are energy, forces, and stress computed at the PBE+U+D3+spin level. D3 is included in the training labels, so no D3 correction is needed at inference time.
• Chemistry: W-S-O systems: 2H-WS₂ and the disordered, oxidized structures produced by O⁺ and O₂⁺ bombardment.
• Energy referencing & standardization: per-element isolated-atom references computed at the same level of theory, combined with tabulated heats of formation; energy, force, and stress targets standardized to the OC20 scale (reusing the OC20 linear-regression coefficients and force RMS).

Usage

This is a state dict to be loaded on top of the base UMA-S 1.1 model using fairchem. You must have access to the gated facebook/UMA repo and be authenticated.

import torch
from fairchem.core import pretrained_mlip, FAIRChemCalculator

device = "cuda"
predictor = pretrained_mlip.get_predict_unit("uma-s-1p1", device=device)

# load fine-tuned weights
ckpt = torch.load("best_checkpoint_conservative.pt", map_location=device)
state = ckpt.get("model_state_dict", ckpt)
cleaned = {}
for k, v in state.items():
    while k.startswith("module."):      # strip DDP/EMA wrapper prefixes
        k = k[len("module."):]
    if k == "n_averaged":               # AveragedModel counter, not a weight
        continue
    cleaned[k] = v
target = getattr(predictor.model, "module", predictor.model)
target.load_state_dict(cleaned)
predictor.model.to(device).eval()

calc = FAIRChemCalculator(predictor, task_name="oc20")
# attach `calc` to an ASE Atoms object as usual

Attribution & citation

This model is fine-tuned from UMA (Universal Models for Atoms) by Meta FAIR Chemistry. Per the FAIR Chemistry License, please acknowledge UMA in any publication that uses this model.

If you use this fine-tuned model, please cite the model itself using the DOI shown in the header of this repository, together with both UMA and our W-S-O paper:

@misc{kwon2026finetuninguniversalmachinelearnedinteratomic,
      title={Fine-Tuning a Universal Machine-Learned Interatomic Potential for Oxygen Plasma Interactions with WS$_2$}, 
      author={Jaehong Kwon and Andrew S. Rosen and David B. Graves},
      year={2026},
      eprint={2606.21632},
      archivePrefix={arXiv},
      primaryClass={cond-mat.mtrl-sci},
      url={https://arxiv.org/abs/2606.21632}, 
}

@article{wood2026family,
  title={UMA: A family of universal models for atoms},
  author={Wood, Brandon and Dzamba, Misko and Fu, Xiang and Gao, Meng and Shuaibi, Muhammed and Barroso-Luque, Luis and Abdelmaqsoud, Kareem and Gharakhanyan, Vahe and Kitchin, John and Levine, Daniel and others},
  journal={Advances in Neural Information Processing Systems},
  volume={38},
  pages={129391--129427},
  year={2026}
}

License

FAIR Chemistry License v1. © Meta Platforms, Inc. and affiliates for the underlying UMA materials. Fine-tuning and derivative weights by the model authors, distributed under the same terms.