Research map / Architectures and efficiency / RNNs and normalization
UMA: A Family of Universal Models for Atoms
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Trains a family of universal atomistic models (UMA, S/M/L) on ~500M atomic systems across materials, molecules, catalysis, molecular crystals and MOFs, using a Mixture of Linear Experts (MoLE) and a two-stage direct-then-conservative training schedule.
Problem
Machine learning interatomic potentials are typically trained per domain on small datasets; a single model that is accurate and fast across diverse DFT tasks (materials, molecules, catalysis, molecular crystals, MOFs) remains an open problem.
Approach
Scale a single eSEN/eSCN-based equivariant MLIP across five merged DFT datasets using Mixture of Linear Experts whose expert weights depend only on global system information (composition, charge, spin, task), so experts can be merged before MD, keeping inference cost near that of a dense model; combined with a BF16-then-FP32 two-stage training recipe and empirical scaling laws.
Key results
- UMA-M achieves the highest Matbench Discovery F1 to date (0.929) among the reported models.
- UMA-L achieves a 25% improvement in AdsorbML success rate over the previous SOTA (EquiformerV2).
- UMA-S simulates 1,000 atoms at 16 steps/s and fits 100k+ atoms on a single 80GB H100; compute-optimal MoLE reaches equivalent loss with ~2.5x fewer active parameters than dense.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use MoLE rather than a dense model for multi-task UMA-S (ablation) | UMA-S · Multi-task training | validation loss | Helped |
| Train multi-task without MoLE rather than single-task (ablation) | small dense models · Multi-task training | validation loss | Hurt |
| Increase MoLE experts from 1 to 8 (ablation) | UMA-S · Multi-task training | loss | Helped |
| Increase MoLE experts from 8 to 32 (ablation) | UMA-S · Multi-task training | loss | Helped |
| Increase MoLE experts from 32 to 128 (ablation) | UMA-S · Multi-task training | loss | No clear effect |
| Use multi-task rather than single-task training (ablation) | UMA-L · Pre-training across tasks | loss | Mixed |
| Use reduced max neighbors in pretraining, then effectively infinite max neighbors in finetuning (ablation) | UMA · Final model performance | performance, energy conservation, and smoothness | No clear effect |
Surprising findings
- BF16-only training degrades MLIP accuracy by as much as 20-50%, unlike LLMs where BF16 is standard.
- The MoLE advantage shrinks at larger scale; a 5.6B-total-parameter MoLE gave only marginal gains over dense at 700M active parameters.
Limitations
- Six-angstrom cutoff can miss long-range interactions beyond 6 Å (e.g., adsorbate starting at 7 Å).
- Charge and spin use discrete embeddings, limiting generalization to unseen charge/spin values.
- UMA-L shows overfitting on the high-entropy alloy test set; materials benchmarks require fine-tuning on MPtrj and sAlex.
Experimental setup
Datasets: OMat24 OMol25 OC20++ OMC25 ODAC25 OC22 OC25 OPoly26 MPtrj sAlex
Benchmarks: Matbench Discovery WBM OC20 S2EF AdsorbML OMol25 OOD-Comp OMol25 PDB-TM OMC25 test CSP Blind Test 7th ODAC OOD MDR phonon elastic tensor NVE MD
Metrics: MAE F1 success rate steps per second Kendall rank lattice energy error
Baselines: eSEN-30M-OMat eqV2-OMat eqV2-OC20 GemNet-OC20 eSEN-sm-cons. eSEN-S-OMC Orb-v3 MACE-MPA-0 SevenNet-MF-ompa GRACE-2L-OAM
Builds on: eSEN eSCN convolution OC20 / OMat24 / OMol25 datasets Chinchilla-style compute-optimal scaling laws
Closest papers by meaning
- UniSim: A Unified Simulator for Time-Coarsened Dynamics of Biomolecules (ICML 2025 · similarity 0.88)
- All-atom Diffusion Transformers: Unified generative modelling of molecules and materials (ICML 2025 · similarity 0.87)
- Benchmarking Compositional Generalisation for Machine Learning Interatomic Potentials (NeurIPS 2026 · similarity 0.86)
- FlashMD: long-stride, universal prediction of molecular dynamics (NeurIPS 2025 · similarity 0.86)
- ChemPile: A 250 GB Diverse and Curated Dataset for Chemical Foundation Models (NeurIPS 2025 · similarity 0.85)
- Beyond Atoms: Enhancing Molecular Pretrained Representations with 3D Space Modeling (ICML 2025 · similarity 0.85)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.