Universal Models for Atoms (UMA-OC20)

• UMA-OC20 is a unified ML framework that leverages the OC20 dataset to model energies, forces, and relaxed geometries across chemically diverse systems.
• It incorporates data augmentation techniques, including rattled structures and short MD trajectories, to enhance robustness and generalization for in-domain and out-of-domain tasks.
• The framework adapts graph-based models such as CGCNN, SchNet, and DimeNet++ for periodic boundary conditions, enabling scalable and accurate predictions in catalysis research.

Universal Models for Atoms (UMA-OC20) refer to the development and application of machine-learned interatomic potentials and representations capable of simultaneously modeling diverse atomistic systems—spanning broad elemental, structural, and chemical variability—using a single unified framework. The OC20 dataset and associated benchmarking tasks formalize this challenge in the field of computational catalysis and surface science, providing the core infrastructure for research on universal atomistic models.

1. Definition and Scope

Universal Models for Atoms (UMA-OC20) are constructed and benchmarked in the context of the Open Catalyst 2020 (OC20) dataset, which provides density functional theory (DFT) calculations of more than 1.2 million structure relaxations, encompassing a chemical space that includes over 55 elements and a range of adsorbates (from simple atomic and molecular species to complex C-, O-, N-containing compounds). These models aim to accurately predict a range of atomic and molecular properties—including but not limited to total energies, atomic forces, and relaxed geometries—for any configuration within this broad compositional and configurational universe, without system-specific retraining or tuning.

This universality is defined both by data diversity—requiring the model to be exposed to, and robust against, chemical, bonding, and structural variability—and by model design, which must enable generalization to out-of-domain (OOD) adsorbates, catalysts, or their combinations.

2. Dataset Construction and Augmentation

The OC20 dataset is the foundational resource underlying UMA-OC20. It consists of 1,281,040 DFT surface relaxations, each accompanied by up to hundreds of single-point evaluations per structure, resulting in approximately 250–265 million total single-point calculations. Surfaces are constructed from stable bulk structures in the Materials Project, and adsorbates include simple O/H species, C₁–C₂ fragments, and N-containing molecules.

To achieve robust model generalization, data augmentation is performed via the inclusion of:

• Off-equilibrium "rattled" structures: Atomic positions are randomly perturbed to sample higher-energy, non-equilibrium configurations.
• Short-timescale ab initio molecular dynamics (MD) trajectories: These add dynamical off-equilibrium states.
• Additional computed properties (e.g., Bader charges, bonding analyses) and periodic boundary condition enforcement.

The dataset is stratified into train/validation/test splits with in-domain and out-of-domain subtasks, supporting rigorous benchmarking of generalization.

3. Machine Learning Model Architectures and Baseline Approaches

Three principal model architectures serve as reference baselines for UMA-OC20:

Model	Core Principle	Notable OC20 Modifications
CGCNN	Graph convolution	Gaussian distance encoding, gradient head
SchNet	Continuous filters	Periodic BCs, force via $\nabla E$
DimeNet++	Directional MP	Periodic BCs, triplet angular encoding

• CGCNN introduces a crystal graph convolutional operation and, for OC20, uses continuous Gaussian basis encodings. Its force prediction is implemented as the negative gradient of the energy head.
• SchNet operates with continuous-filter convolutions; for this problem, it is extended to periodic boundary conditions and a composite loss for energies and forces.
• DimeNet++ explicitly includes triplet angular correlations in message passing, adapted for periodicity and force outputs, and further improved upon scaling.

All models use graph-based constructions: atoms are nodes, edges connect all atoms within a 6 Å cutoff (including relevant periodic images). This graph and associated features encode the local chemical environment necessary for universal description.

4. Benchmark Tasks and Evaluation Framework

OC20 defines three key benchmark tasks central to catalyst modeling workflows:

1. S2EF (Structure → Energy and Forces): Given atomic coordinates, predict both adsorption energies ($E_{ad}$) and per-atom forces. Adsorption energies are computed as $E_{ad} = E_{sys} - E_{slab} - E_{gas}$ with gas-phase reference energies obtained via linear atomic decompositions and tabulations.
2. IS2RS (Initial Structure → Relaxed Structure): Predict the DFT-relaxed geometry from an unrelaxed initial state.
3. IS2RE (Initial Structure → Relaxed Energy): Directly predict the energy of the ultimately relaxed geometry, bypassing the relaxation procedure.

Each task has an associated suite of metrics (energy MAE, force MAE, force cosine similarity, etc.), and evaluation is performed on both in-domain and out-of-domain splits (e.g., unseen adsorbates, unseen catalysts, or both). The benchmarking infrastructure includes public leaderboards, open software (PyTorch Geometric-based), and comprehensive documentation.

5. Model Scaling, Generalization, and Extrapolation

A key empirical result is the absence of clear saturation with respect to model size across all evaluated metrics and tasks. For example, increasing parameter count in SchNet or employing larger DimeNet++ variants continues to improve performance. Notably, including off-equilibrium data is statistically beneficial for force cosine similarity and downstream relaxation tasks (IS2RS).

Research shows modest differences between in-domain and out-of-domain metrics, indicating that larger models with more diverse data may eventually achieve full universality—predictive performance with negligible chemical or structural specificity. Incorporating MD and rattled configurations in training improves the quantity and diversity of the force signal, enhancing generalization across bonding environments and chemistries.

6. Technical and Computational Details

Key technical elements specific to UMA-OC20 models and their training include:

• Composite energy/force loss: $$\mathcal{L} = \lambda_E \sum |E_i - E_i^{DFT}| + \lambda_F \sum \frac{1}{N_i} \sum_j |F_{i,j} - F_{i,j}^{DFT}|$$, where forces are computed as $F = -\nabla_E$.
• Graph construction: Radius-based graphs augmented with periodic images for proper local/periodic environment encoding.
• Hyperparameter scaling: Systematic benchmarking of embedding sizes, cutoff radii, Gaussian widths, and number of message-passing layers elucidates the relationship between data/model size and achievable accuracy.
• Adsorption energy referencing: Gas-phase reference values are linearly decomposed and tabulated.

These choices enable robust model development across diverse system sizes (from hundreds of thousands to over 100 million samples) and contribute to the observed scalability.

7. Open Science and Community Infrastructure

The OC20 initiative is structured to facilitate widespread adoption and continual improvement of UMA-OC20 approaches:

• Full release of the OC20 dataset, including all DFT trajectories and analysis outputs.
• Open-source codebase for data loaders, preprocessing, and training.
• A public leaderboard and standardized split definitions allow the global research community to compare models on consistent, reproducible tasks.
• Supplementary resources (e.g., electronic structure analyses, curated gas-phase reference tables) and documentation lower the barrier for new entrants to the field and promote rigorous comparative studies.

Such infrastructure is intended to transition computational catalysis and related fields towards the routine use of universal, foundation-level models for atomistic simulation workflows.

---

In summary, Universal Models for Atoms (UMA-OC20), grounded in the OC20 dataset and challenge framework, establish a paradigm in which atomistic machine learning potentials are trained to deliver quantitatively accurate, generalizable predictions across broad classes of chemical systems. The combination of dataset breadth, model scalability, open benchmarking, and community engagement constitutes the essential foundation for ongoing advances in universal atomistic modeling for catalysis and materials science (Chanussot et al., 2020).