UMA: A Family of Universal Models for Atoms

Brandon M. Wood; Misko Dzamba; Xiang Fu; Meng Gao; Muhammed Shuaibi; Luis Barroso-Luque; Kareem Abdelmaqsoud; Vahe Gharakhanyan; John R. Kitchin; Daniel S. Levine; Kyle Michel; Anuroop Sriram; Taco Cohen; Abhishek Das; Ammar Rizvi; Sushree Jagriti Sahoo; Zachary W. Ulissi; C. Lawrence Zitnick

Source-linked AI summary

UMA: A Family of Universal Models for Atoms

Brandon M. Wood, Misko Dzamba, Xiang Fu, Meng Gao, Muhammed Shuaibi, Luis Barroso-Luque, Kareem Abdelmaqsoud, Vahe Gharakhanyan, John R. Kitchin, Daniel S. Levine, Kyle Michel, Anuroop Sriram, Taco Cohen, Abhishek Das, Ammar Rizvi, Sushree Jagriti Sahoo, Zachary W. Ulissi, C. Lawrence Zitnick

arXiv:2506.23971v2Jun 30, 2025cs.LG

TL;DR

Computationally expensive DFT motivates MLIPs that can generalize across chemistry and materials tasks. UMA trains a family of models on nearly half a billion atomic systems, using scaling laws, Mixture of Linear Experts, and two-stage training. Without fine-tuning, UMA performs similarly or better than specialized models across broad benchmarks, supporting practical research applications while retaining stated misuse concerns.

Problem
DFT is computationally expensive, and MLIPs that generalize across chemical domains and specialized DFT tasks remain difficult to develop.
Method
UMA compiles ∼500M atomic systems across domains and uses empirical scaling laws, Mixture of Linear Experts, and a two-stage training schedule.
Results
Without fine-tuning, UMA performs similarly or better than specialized models in accuracy and inference-speed/memory-efficiency across material, molecular, and catalysis benchmarks.
Takeaways & Limitations
A single model can achieve sufficient accuracy for practical research applications across a broad spectrum of chemistry and materials science applications.
Takeaways & Limitations
Potential misuses exist, and the authors mitigate these risks by training on datasets designed for applications beneficial to society.

Abstract

from arXiv · show

The ability to quickly and accurately compute properties from atomic simulations is critical for advancing a large number of applications in chemistry and materials science including drug discovery, energy storage, and semiconductor manufacturing. To address this need, Meta FAIR presents a family of Universal Models for Atoms (UMA), designed to push the frontier of speed, accuracy, and generalization. UMA models are trained on half a billion unique 3D atomic structures (the largest training runs to date) by compiling data across multiple chemical domains, e.g. molecules, materials, and catalysts. We develop empirical scaling laws to help understand how to increase model capacity alongside dataset size to achieve the best accuracy. The UMA small and medium models utilize a novel architectural design we refer to as mixture of linear experts that enables increasing model capacity without sacrificing speed. For example, UMA-medium has 1.4B parameters but only ~50M active parameters per atomic structure. We evaluate UMA models on a diverse set of applications across multiple domains and find that, remarkably, a single model without any fine-tuning can perform similarly or better than specialized models. We are releasing the UMA code, weights, and associated data to accelerate computational workflows and enable the community to continue to build increasingly capable AI models.

1 Introduction

DFT is valuable across chemistry and materials science but too computationally expensive for broad use, motivating faster MLIPs that generalize across tasks. UMA combines nearly half a billion atomic systems to test whether one model can achieve strong accuracy, speed, and generalization across domains.

DFT supports applications from drug discovery to semiconductors, but its computational expense limits usage.
MLIPs can approximate DFT with lower computational complexity, but generalization across chemical domains and specialized tasks remains an open problem.The passage contrasts O(n) MLIPs with O(n3) DFT and defines tasks by domain plus DFT settings.
Pooling datasets across tasks enables larger multi-task models, addressing the difficulty of scaling atomic datasets because of their computational cost.
UMA uses ∼500M atomic systems and empirical scaling laws to study accuracy, speed, and generalization for a single model across chemistry and materials science.The combined datasets cover nearly the entire chemical space except radioactive elements.
Without task-specific fine-tuning, UMA performs similarly or better than specialized models across material, molecular, and catalysis benchmarks.Reported highlights include state-of-the-art Matbench Discovery results and a 25% improvement in successful catalysis adsorption-energy calculations.

2 Approach

UMA extends an equivariant atomic model into a family spanning accuracy, speed, and use cases, while incorporating task information and Mixture of Linear Experts for efficient multi-task scaling. Its routing uses global system information, enabling precomputation, and training combines direct-force pretraining with conserving-force and stress fine-tuning.

UMA model family: UMA models estimate atomic energy, forces, and optionally stress from atomic inputs, supporting simulations that require thousands or millions of force evaluations.UMA-S targets computationally intensive long molecular dynamics, while UMA-L is intended as a highly accurate DFT surrogate.
UMA model family: The architecture extends eSEN with charge, spin multiplicity, and DFT-task inputs so one model can represent multiple chemical settings.These inputs are embedded and added to node representations across message-passing layers.
Mixture of Linear Experts: MoLE replaces linear operations with mixtures of linear experts, increasing capacity while preserving rotational equivariance and smooth variation across tasks.Each expert has weights W_k and contribution α_k, with dense expert use rather than sparse competition.
Mixture of Linear Experts: Routing weights depend on global composition, charge, spin, and task information, allowing effective weights to be precomputed before long simulations.Position and neighborhood information are excluded because they can vary during relaxations or molecular dynamics.
Training: The training pipeline first predicts forces directly, then removes the force head and fine-tunes with automatic differentiation to predict conserving forces and stresses.The two-stage procedure addresses the resource demands of training conservative models.

3 Datasets

UMA pools molecular, materials, catalytic, interaction, and molecular-crystal datasets into a nearly half-billion-example training collection spanning diverse chemical environments. The paper evaluates generalization through held-out domain test sets and practical benchmarks, while scaling analyses compare dense and MoLE architectures.

Training datasets: OMol25, OMat24, OC20, ODAC25, and OMC25 together provide close to 500 million training examples and over 30 billion atoms.The combined datasets cover materials, molecules, molecule–surface interactions, and molecular crystals.
Training datasets: The datasets use different DFT settings, including PBE/VASP for OMat24 and ωB97M-V/ORCA for OMol25, requiring a unified model to handle heterogeneous supervision.Sampling ratios are adjusted to 4 for OMat24 and OMol25, 1 for OC20 and ODAC25, and 2 for OMC25.
Scaling analysis: Figure 3 measures validation loss against FLOPs, model size, and dataset size for dense and 8-expert MoLE models, identifying compute-optimal frontiers.At equivalent loss, a compute-optimal MoLE model can use approximately 2.5× fewer active parameters.
Evaluation: Table 2 reports MAE on held-out materials, catalysis, molecules, molecular-crystal, and MOF test splits, with energies, forces, and stresses shown in distinct units.UMA results are compared with state-of-the-art literature models and approximate practical-utility targets.
Evaluation: Figure 4 compares normalized UMA-L pretraining errors for single-task and multi-task models, showing frequent single-task force overfitting and generally lower multi-task convergence errors.The comparison uses models with the same architecture and size.

4 Results

UMA demonstrates scaling behavior and broad multi-domain performance, with mixture of linear experts improving parameter efficiency while preserving accuracy and inference practicality.

Model and Data Scaling: UMA models exhibit log-linear scaling behavior across the tested 10^18–10^20 FLOP range, indicating that greater capacity is needed for the dataset.The scaling relationships guided model-size selection and showed an advantage for MoLE over dense architectures.
Model and Data Scaling: An optimal MoLE model can match the loss of a dense model approximately 2.5 ± 0.2 times larger for UMA-M.The advantage decreases at larger model sizes, where dense and MoLE performance converge.
Multi-task vs. Single-Task: MoLE enables UMA-S multi-task training to achieve results comparable to task-specialized small models, whereas dense multi-task models perform worse than single-task models in this regime.Increasing UMA-S from 1 to 8 experts produced a significant loss improvement, with a smaller gain at 32 experts.
Multi-task vs. Single-Task: Multi-task training benefits large models even without MoLE, while single-task models frequently overfit forces during further training.The exception is OMol forces, where both models already achieve errors below 10 meV/Å.
Inference Efficiency: MoLE weight merging avoids additional inference-time or memory penalties for long-timescale simulations on a single system.This supports larger length and time scales for molecular dynamics, relaxations, and related simulations.
Evaluation: UMA performs similarly or better than specialized models across material, molecular, and catalysis benchmarks without task-specific fine-tuning.The evaluation covers held-out test splits and practically important benchmarks across materials, catalysis, molecules, molecular crystals, and MOFs.
Evaluation: UMA-M achieves the highest F1 score to date on Matbench Discovery, while UMA models reduce OC20 S2EF adsorption-energy errors by around 80%.UMA also performs strongly on molecular, molecular-crystal, and MOF evaluations, including the hardest ODAC OOD test set.

5 Related Work

Universal MLIPs seek zero-shot generalization across chemical domains, but domain shifts and differing DFT settings make this difficult; UMA studies pooled data and scaling relationships as routes forward.

Universal MLIPs: Materials-focused MLIPs called universal may not generalize well to molecules or surfaces because chemical, structural, elemental, and DFT-theory distributions differ.For example, PBE is used for materials while ωB97M-V is used for molecules.
Universal MLIPs: High-accuracy zero-shot performance across multiple domains remains an open challenge, despite progress from large pretraining followed by task-specific fine-tuning.Fine-tuned models have performed significantly better than models trained from scratch, but removing specialization would increase usefulness.
Scaling Laws: Empirical scaling relations help predict model and dataset sizes and allocate compute, but their behavior for very large MLIP training runs had remained uncertain.UMA reports log-linear scaling in the tested FLOP range and uses it to select model sizes and compare MoLE with dense architectures.

6 Limitations

UMA’s principal limitations concern long-range interactions and the representation of charge and spin, while the authors also acknowledge potential misuse risks.

Long-range interactions: The 6Å cutoff in UMA-S and UMA-M can mishandle molecules separated from surfaces by more than 6Å.An adsorbate beginning 7Å from a catalyst surface is viewed as an independent non-interacting structure.
Charge and spin: Separate embeddings for each discrete charge or spin limit generalization to unseen charges or spins.The authors identify improved charge and spin incorporation as an area for future work.
Potential misuse: The authors note that potential misuses exist and describe training on datasets designed for beneficial societal applications as a mitigation measure.

7 Discussion and Conclusion

UMA trains MLIPs across diverse DFT tasks using nearly 500 million examples, combining scaling-law analysis with Mixture of Linear Experts to balance accuracy and inference efficiency. Across materials, molecules, catalysts, molecular crystals, and MOFs, UMA achieves strong benchmark performance and supports practical simulations.

Nearly 500 million training examples span numerous datasets and diverse DFT tasks.
Mixture of Linear Experts increases model capacity while maintaining inference efficiency.
UMA-S offers a speed–accuracy balance and simulates 1,000 atoms at 1.4ns per day on one 80GB GPU.
UMA demonstrates strong performance across materials, molecules, catalysts, molecular crystals, and metal organic frameworks.
The two-stage training procedure first predicts forces directly, then fine-tunes for conserving forces and stresses using auto-grad.
A heat-of-formation reference makes energy magnitudes comparable across datasets with different DFT settings.

A.7 UMA 1.2 model and training

UMA 1.2 introduces training and architectural changes including conservative FP32 training, charge balancing, task-specific output heads, and expanded data handling. The release documentation indicates that updated code and configuration files will accompany the release.

UMA 1.2 uses conservative FP32 training end-to-end.
The model adds charge balancing to L0 channels and task-specific output heads.
UMA 1.2 enables single-atom predictions and includes single-atom data for all datasets.
OMat24 is separated into AIMD and Rattled subsets with different data ratios and loss coefficients.
The updated code and configuration files will be available in the fairchem repository at release.

B Training Data

UMA’s training data combines large, diverse datasets across inorganic materials, molecules, catalysis, molecular crystals, and MOFs. The scaling analysis accounts for the computational cost of equivariant atomic networks and supports compute-aware model design.

459 million training examples across five datasets contain systems with up to 350 atoms.
OMat24 contributes 100M diverse inorganic-material structures calculated with VASP and PBE.
B.2 Molecules: OMol25 spans 80+ elements and includes metal complexes, biomolecules, electrolytes, and existing datasets under one theory level.
The Open Molecular Crystals dataset contributes 25 million molecular crystal structures to address scarce public data.
Scaling-law experiments use lmax = 4, mmax = 2, and nneighbors = 30, yielding approximately 270 FLOPs/parameter/atom per training step.
Equivariant-GNNs require 2–3 orders of magnitude more FLOPs than similarly sized LLM networks because of higher parameter reuse.

C.2 Compute optimal fits

The paper fits compute-optimal model and dataset sizes with power-law relationships, then benchmarks inference and predictive performance across materials, catalysis, molecules, and related tasks. The evaluations combine broad benchmark coverage with implementation-specific efficiency measurements.

C.2 Compute optimal fits: Compute-optimal model size N∗(C) and dataset size D∗(C) are fitted as power laws of training compute C.
C.2 Compute optimal fits: Two methods for estimating the loss relationship produced similar results, while direct minimization was more sensitive to hyperparameters.
Pre-merged MoLE experts and activation checkpointing allow UMA-S to fit more than 100k atoms into memory.
Non-MoLE merging has comparable inference speed but requires more GPU memory to store parameters.
Table 14 covers Matbench-Discovery, MDR phonon, elastic tensor, AdsorbML, molecule, molecular-crystal, and NVE MD benchmarks.
For catalysis, adsorption energies are computed by subtracting two total-energy calculations, improving results over prior models.

E.4 Molecules

UMA models are evaluated across molecular and molecular-crystal tasks, where larger models and cross-domain training generally improve performance. UMA also outperforms task-specific baselines on molecular-crystal evaluations.

Molecules: UMA-S and UMA-S-OMol provide comparable results on OMol25 evaluations using the preview dataset.The comparison covers validation, test, single-point, and optimization tasks.
Molecules: UMA-M and UMA-L perform significantly better than UMA-S and UMA-S-OMol on OMol25 single-point and optimization evaluations.The reported explanation is that UMA-M is energy conserving and better behaved during optimization tasks.
Molecular Crystals: All UMA sizes outperform the UMA-S-OMC model trained only on OMC25, indicating complementary information from datasets such as OMol25.The evaluation concerns molecular-crystal packing and requires accurate intermolecular forces.
Molecular Crystals: Using total energies produces significantly better OpenDAC adsorption-energy estimates, while UMA-M and UMA-L have forces similar to SoTA eqV2-ODAC.OpenDAC contains metal-organic frameworks with CO2 and water molecules.
Molecular Crystals: Expert usage overlaps strongly between OC20 and OMat24, whereas OMol25 has minimal overlap and OMC25 and ODAC23 share one expert.OMC25 and ODAC23 utilize the fewest experts among the analyzed datasets.

F.2 Generalization Across Architectures

MoLE layers are tested across two distinct MLIP architectures under controlled conditions. Their relative benefit is reported as consistent across architectures, suggesting a general capacity increase without inference speed or memory overhead.

F.2 Generalization Across Architectures: 8-expert MoLE layers are evaluated in eSEN and EquiformerV2 using 6M base parameters and 8 Omol training epochs.Both experiments use direct pretraining only and control for model size and training duration.
F.2 Generalization Across Architectures: MoLE provides a fairly consistent relative benefit in eSEN and EquiformerV2 despite their architectural differences.The comparison measures relative performance improvements over versions without MoLE.
F.2 Generalization Across Architectures: MoLE is presented as a general approach for increasing MLIP capacity without inference-speed or memory overhead.The conclusion follows the cross-architecture comparison of eSEN and EquiformerV2.

G Diatomics

Diatomic curves expose weaknesses in earlier UMA versions and show that updated training produces smooth predictions. Adding diatomic data substantially improves accuracy, but OMol asymptotic issues remain due to global charge/spin handling and the cutoff envelope.

G Diatomics: Earlier UMA versions showed unphysical diatomic behavior, including large OMat energy fluctuations beyond the bonding region.The source of these issues was unclear among the model, data, and training procedure.
G Diatomics: Updated UMA-1.2 produces qualitatively smooth OMat and OMol diatomic curves even without diatomic training data.This result concerns the unreleased UMA-1.2 version trained without added diatomic data.
G Diatomics: Adding OMat and OMol diatomic data significantly improves diatomic energy and force estimates, and the resulting model is released publicly.The two UMA-1.2 versions differ only by the additional diatomic data.
G Diatomics: OMol asymptotic behavior remains problematic because global charge and spin interact with an envelope function that decays interactions to zero at the cutoff.At approximately 5.5 Å, the model treats the separated atoms as disconnected while retaining the system’s global charge and spin.