Fine-Tuning a Universal Machine-Learned Interatomic Potential for Oxygen Plasma Interactions with WS$_2$

Technical Summary: Fine-Tuning a Universal Machine-Learned Interatomic Potential for Oxygen Plasma Interactions with WS2

Problem Statement
Molecular dynamics (MD) simulations of plasma-surface interactions require interatomic potentials (IPs) that balance accuracy, computational efficiency, and the ability to describe complex, reactive systems involving multiple elements and bonding types. Traditional IPs often struggle to meet these demands simultaneously, particularly for industrially relevant scenarios like the atomic layer etching (ALE) of two-dimensional transition metal dichalcogenides (TMDs) such as tungsten disulfide (WS2). A specific challenge arises in modeling oxygen plasma interactions with WS2, where the chemistry involves spin-polarized oxygen species (triplet O2) and requires Hubbard U corrections to accurately describe the W 5d states in oxidized tungsten (WOx). Standard foundation models, while powerful, are often trained on datasets (e.g., Open Catalyst 2020) that lack high-energy bombardment configurations and specific electronic corrections (spin polarization, U) necessary for this chemistry.

Methodology
The authors employ an iterative fine-tuning loop to adapt the Universal Models for Atoms (UMA) foundation model, specifically the uma-s-1p1 variant trained on the Open Catalyst 2020 (OC20) task, for oxygen plasma interactions with multilayer WS2.

1. Simulation Setup: MD simulations model the bombardment of a 2H-WS2 slab by 15 eV O+ and O2+ ions. The system treats incoming ions as fast neutrals (due to rapid charge transfer) and utilizes a stratified Berendsen thermostat to manage post-impact heating. Detached clusters (SO, SO2, O, O2) are dynamically identified and removed based on specific geometric and compositional criteria.
2. Iterative Fine-Tuning Loop:
   • MD Exploration: The current model generates trajectories to explore the W-S-O configuration space, utilizing both flat and stepped slab geometries.
   • Configuration Sampling: A maximally diverse subset of configurations is selected using Smooth Overlap of Atomic Positions (SOAP) descriptors, Principal Component Analysis (PCA), and Farthest Point Sampling (FPS).
   • DFT Labeling: Selected configurations are labeled using Density Functional Theory (DFT) at the PBE+D3+U+spin level. This includes spin polarization and a Hubbard U correction (2.87 eV on W 5d orbitals) to address the electronic structure limitations of the pretrained OC20 model.
   • Model Update: The model is fine-tuned on the new labeled data. The authors primarily utilize a layer freezing strategy, updating only selected layers while keeping the majority of the model parameters frozen, though full fine-tuning and Low-Rank Adaptation (LoRA) were also tested for comparison.
3. Evaluation: A held-out test set is generated using the updated model to ensure unbiased evaluation of generalization capabilities. Convergence is assessed via Mean Absolute Error (MAE) on energy, force, and stress, as well as the convergence of physical observables (chemisorbed S and O coverage).

Key Contributions

• Adaptation of Foundation Models: The study demonstrates the viability of fine-tuning a general-purpose foundation model (UMA) for a specific, high-energy plasma-surface application involving complex electronic corrections (spin and U) not present in the original training data.
• Iterative Sampling Strategy: The implementation of an iterative loop combining SOAP-based diversity sampling with DFT labeling ensures the training set covers the relevant configuration space, including high-energy impact states and varying oxidation levels.
• Efficient Training Strategy: The work highlights the efficacy of layer freezing in fine-tuning large foundation models, achieving high accuracy with a reduced number of trainable parameters compared to full fine-tuning.

Results

• Accuracy Improvements: After three rounds of fine-tuning (R3), the model achieved a Mean Absolute Error (MAE) of 4.5 × 10⁻³ eV/atom for energy and 0.076 eV/Å for force on a held-out test set. This represents a significant improvement over the pretrained baseline (evaluated against the same PBE+D3+U+spin reference), which had a force MAE of approximately 0.23 eV/Å.
• Error Distribution: The fine-tuning process not only reduced the average error but also tightened the distribution of errors, reducing the 99% inclusion bound for force errors from ±1.343 eV/Å (pretrained) to ±0.479 eV/Å (R3), thereby enhancing MD trajectory stability.
• Physical Observables: Crucially, the study found that the pretrained UMA-S model (without fine-tuning) already reproduced the key production-scale observables: the monotonic removal of chemisorbed sulfur and the uptake of chemisorbed oxygen to form WOx. The fine-tuned models (R2 and R3) converged to the same physical trajectories as the pretrained model, despite the latter lacking spin polarization and Hubbard U corrections. The R1 model showed slight deviations, consistent with its higher error metrics.
• Local Environment Analysis: BIRCH clustering analysis revealed that the largest initial discrepancies occurred in high W-O coordination environments, which were effectively corrected by the third round of fine-tuning.

Significance and Claims
The paper claims that foundation models like UMA can serve as a strong baseline for plasma-surface simulations, even for chemistries involving specific electronic corrections (spin, U) that were absent in the pretraining data. The primary significance lies in demonstrating that fine-tuning significantly improves the underlying potential energy surface accuracy (reducing MAE and stabilizing error distributions) while the production-scale observables of interest (S removal and O uptake) were already well-captured by the pretrained model. This suggests that for certain plasma-surface applications, the computational cost of extensive fine-tuning might be weighed against the fact that the foundation model already provides a qualitatively correct description of the macroscopic dynamics, though fine-tuning remains essential for achieving quantitative DFT-level accuracy in the underlying forces and energies. The authors note that a quantitative comparison with experimental measurements is reserved for a future publication.