UMA: A Family of Universal Models for Atoms

Meta FAIR introduces UMA, a family of universal models for atoms trained on half a billion 3D structures using a novel mixture of linear experts architecture, which achieves state-of-the-art speed and accuracy across diverse chemical domains without requiring fine-tuning.

The Gist

Imagine you are trying to build a house. To do this, you need to know exactly how every brick, beam, and nail interacts with the others. In the world of chemistry and materials science, these "bricks" are atoms. For decades, scientists have used a super-precise but incredibly slow computer program called Density Functional Theory (DFT) to figure out how atoms behave. It's like using a master architect to draw every single blueprint by hand. It's accurate, but it takes hours or even days to simulate just a tiny interaction. This slows down everything from discovering new life-saving drugs to creating better batteries.

Enter UMA (Universal Models for Atoms), a new family of AI models from Meta FAIR that acts like a "super-fast, super-smart apprentice" who can predict how atoms will behave in a fraction of a second, with almost the same accuracy as the master architect.

Here is the simple breakdown of how they did it:

1. The "Giant Library" of Atomic Experiences

To teach an AI to be good at something, you usually need to show it a lot of examples. Previous AI models for atoms were like students who only studied one subject (like only learning about metals, or only learning about water molecules). If you asked them about something outside their specific subject, they got confused.

The UMA team decided to build a giant, universal library. They gathered data on half a billion unique 3D atomic structures. This includes:

• Materials: Like the stuff in your phone or car.
• Molecules: Like the chemicals in medicine.
• Catalysts: The "helpers" that speed up chemical reactions.
• Crystals and Frameworks: Complex structures used in air filtration.

By feeding the AI this massive, diverse diet, they taught it to understand the "language" of atoms in general, rather than just memorizing specific recipes.

2. The "Mixture of Linear Experts" (MoLE) – The Swiss Army Knife

Usually, to make an AI smarter, you just make it bigger. But making a model bigger usually makes it slower. It's like trying to carry a giant library in your backpack; the more books you add, the harder it is to walk.

The UMA team invented a clever trick called Mixture of Linear Experts (MoLE).

• The Analogy: Imagine a massive team of 32 specialized mechanics (experts). When a car (an atomic structure) comes in, the team leader doesn't ask all 32 mechanics to work on it. Instead, they ask only the 3 or 4 mechanics who are best suited for that specific car.
• The Result: The model has a huge "brain" (1.4 billion parameters, like having a massive library of knowledge), but for any single calculation, it only "wakes up" a tiny fraction of that brain (about 50 million parameters).
• Why it matters: This means UMA is incredibly smart (it knows about almost everything) but runs just as fast as a much smaller, simpler model. It's like having a supercomputer in your pocket that doesn't drain your battery.

3. The "Two-Stage Training" – Learning to Walk Before Running

Training these models is hard because you don't want them to just guess; you want them to follow the laws of physics (like energy conservation).

• Stage 1: They taught the AI to predict forces quickly and roughly. Think of this as teaching a child to run without worrying about tripping.
• Stage 2: They refined the model to ensure it follows the strict rules of physics perfectly, ensuring that if you simulate a chemical reaction, the energy doesn't magically appear or disappear.

4. Why This Changes Everything

The paper shows that UMA is a "one-size-fits-all" solution that actually works better than specialized tools.

• Speed: It can simulate 1,000 atoms moving around in real-time on a single computer chip. This is a massive leap forward.
• Accuracy: It beats the current best models in predicting how new materials will hold up, how drugs will bind to viruses, and how to design better catalysts for green energy.
• No Fine-Tuning Needed: Usually, if you want an AI to do a specific job (like designing a battery), you have to re-train it specifically for that. UMA is so well-trained on the "Giant Library" that you can just ask it to do the job, and it does it immediately without extra training.

The Bottom Line

Think of UMA as the Google Maps of the atomic world. Before, if you wanted to know the best route for a specific type of car, you had to hire a different expert for every car model. Now, with UMA, you have one universal map that knows the terrain for every type of car, vehicle, and terrain, and it gives you the fastest route instantly.

This technology is being released to the public to help scientists accelerate the discovery of new medicines, cleaner energy, and better materials, turning what used to take years of calculation into seconds.

Technical Summary

1. Problem Statement

Density Functional Theory (DFT) is the gold standard for simulating atomic interactions in chemistry and materials science, underpinning applications in drug discovery, energy storage, and semiconductor manufacturing. However, DFT is computationally expensive ($O(n^3)$ complexity), limiting its use in large-scale simulations like molecular dynamics (MD) or high-throughput screening.

While Machine Learning Interatomic Potentials (MLIPs) offer a faster alternative ($O(n)$), current models face significant limitations:

• Specialization: Most MLIPs are trained on specific domains (e.g., only materials or only molecules) and fail to generalize across different chemical environments or DFT settings.
• Scaling Trade-offs: Increasing model size to improve accuracy often leads to prohibitive inference costs, making them unsuitable for long-timescale simulations.
• Data Fragmentation: Training on diverse datasets with different DFT settings (functionals, codes like VASP vs. ORCA) is difficult due to energy referencing issues and distribution shifts.

2. Methodology

The authors introduce UMA (Universal Models for Atoms), a family of models designed to be a single, general-purpose surrogate for DFT across diverse chemical domains.

A. Data Strategy

UMA is trained on a massive, unified dataset comprising ~500 million unique 3D atomic structures (over 30 billion atoms), compiled from five major open datasets:

• OMat24: Inorganic materials.
• OMol25: Organic and inorganic molecules.
• OC20/OC22/OC25: Catalysis (adsorbate-surface interactions).
• OMC25: Molecular crystals.
• ODAC25: Metal-Organic Frameworks (MOFs) for direct air capture.

To handle the heterogeneity of DFT settings (e.g., PBE vs. $\omega$B97M-V functionals), the authors employ:

• Task Embeddings: A global vector representing the specific DFT settings (functional, code, charge, spin) is concatenated with atom embeddings.
• Energy Referencing: A "Heat of Formation" (HOF) referencing scheme normalizes energy labels across datasets to ensure comparable numerical ranges.

B. Architecture: Mixture of Linear Experts (MoLE)

To balance high capacity with inference speed, UMA utilizes a novel Mixture of Linear Experts (MoLE) architecture, distinct from the sparse, non-linear MoEs common in Large Language Models (LLMs).

• Linear Experts: The model replaces standard linear layers with a set of linear experts ($y = \sum \alpha_k W_k x$).
• Dense Routing: Unlike sparse MoEs, UMA encourages the dense use of all experts to ensure smooth energy surfaces and rotational equivariance (crucial for physics).
• Pre-computation: Crucially, the routing weights ($\alpha_k$) depend only on time-invariant global information (element composition, charge, spin, task). This allows the model to pre-compute a merged weight matrix ($W^* = \sum \alpha_k W_k$) before simulation.
  • Result: UMA achieves the capacity of a massive model (e.g., 1.4B total parameters) with the inference speed of a much smaller model (e.g., ~50M active parameters), enabling efficient long-timescale MD.

C. Training Procedure

• Two-Stage Training:
  1. Direct Pre-training: The model predicts forces directly using BF16 precision for speed and stability.
  2. Conservative Fine-tuning: The force head is removed, and the model is fine-tuned using automatic differentiation (autograd) to ensure energy conservation and smooth potential energy landscapes. This stage uses FP32 precision to recover accuracy lost in BF16.
• Scaling Laws: The authors derived empirical scaling laws relating compute, data, and model size to determine the optimal architecture for the 500M-structure dataset.

3. Key Contributions

1. Universal Generalization: Demonstrated that a single model, without fine-tuning to specific tasks, can achieve state-of-the-art (SOTA) or comparable performance across materials, molecules, catalysis, molecular crystals, and MOFs.
2. MoLE Architecture: Introduced a physics-aware Mixture of Linear Experts that scales model capacity without sacrificing inference speed, solving the "accuracy vs. efficiency" bottleneck in MLIPs.
3. Empirical Scaling Laws for MLIPs: Established that log-linear scaling relationships (similar to LLMs) hold for atomic models, guiding the optimal allocation of compute resources between model size and dataset size.
4. Massive Unified Dataset: Curated and released a training set of ~500M atomic systems, covering nearly the entire chemical space (excluding radioactive elements).

4. Results

The UMA family includes three model sizes: UMA-S (Small), UMA-M (Medium), and UMA-L (Large).

• Accuracy:
  • Materials: UMA-M achieves the highest F1 score on the Matbench Discovery leaderboard.
  • Catalysis: UMA models reduce adsorption energy errors by ~80% on the OC20 benchmark compared to previous SOTA. On the AdsorbML benchmark, UMA-L improves success rates by 25%.
  • Molecules: UMA-M achieves ligand-strain errors approaching DFT reference levels, making it viable for structure-based drug design.
  • Crystals: UMA-S outperforms specialized baselines in predicting lattice energies and ranking polymorphs in the 7th CSP Blind Test.
• Efficiency:
  • UMA-M (1.4B total params, 50M active) runs at 3 steps/second for 1,000 atoms on an H100 GPU.
  • UMA-S can simulate 100,000+ atoms in memory on a single 80GB GPU, running at 1.4 ns/day for 1,000 atoms.
  • Unlike dense models, UMA does not suffer from memory overflow (OOM) at large system sizes due to the MoLE pre-merging strategy.
• Generalization: UMA models show robust performance on Out-of-Distribution (OOD) tests (e.g., high-entropy alloys, unseen elements) without task-specific fine-tuning.

5. Significance

The UMA models represent a paradigm shift in computational chemistry and materials science:

• Democratization of High-Fidelity Simulation: By providing a single, fast, and accurate model that works across domains, UMA removes the need for researchers to select and fine-tune specialized models for every new problem.
• Enabling New Science: The ability to run accurate, energy-conserving simulations on systems with 100,000+ atoms and long time scales unlocks research into phenomena previously inaccessible to DFT or specialized MLIPs (e.g., complex defect dynamics, large-scale biomolecular interactions).
• Open Science: The authors are releasing the code, weights, and training data under a commercially permissive license, accelerating the development of the next generation of AI-driven scientific discovery.

In summary, UMA successfully bridges the gap between the accuracy of DFT and the speed of machine learning, establishing a new standard for universal, general-purpose atomic simulation.